feat: enhance contour builder with style parameters

Author: PierreRaybaut

plotpy/builder/image.py

@@ -826,6 +826,9 @@ def contours(
         levels: float | np.ndarray,
         X: np.ndarray | None = None,
         Y: np.ndarray | None = None,
+        color: str | None = None,
+        linestyle: str | None = None,
+        linewidth: float | None = None,
     ) -> list[ContourItem]:
         """Make a contour curves
 
@@ -845,11 +848,24 @@ def contours(
              ``numpy.meshgrid``), or it must both be 1-D such that ``len(Y) == N``
              is the number of rows in *Z*.
              If none, they are assumed to be integer indices, i.e. ``Y = range(N)``.
+            color: contour color name. Default is None
+            linestyle: contour line style (MATLAB-like string or "SolidLine",
+             "DashLine", "DotLine", "DashDotLine", "DashDotDotLine", "NoPen").
+             Default is None
+            linewidth: contour line width (pixels). Default is None
 
         Returns:
             list of :py:class:`.ContourItem` objects
         """
-        return create_contour_items(Z, levels, X, Y)
+        return create_contour_items(
+            Z,
+            levels,
+            X,
+            Y,
+            color=color,
+            linestyle=linestyle,
+            linewidth=linewidth,
+        )
 
     def histogram2D(
         self,

plotpy/items/contour.py

@@ -29,7 +29,7 @@
 
 from plotpy.config import _
 from plotpy.items.shape.polygon import PolygonShape
-from plotpy.styles import ShapeParam
+from plotpy.styles import COLORS, ShapeParam
 
 
 class ContourLine(gds.DataSet):
@@ -123,6 +123,9 @@ def create_contour_items(
     levels: float | np.ndarray,
     X: np.ndarray | None = None,
     Y: np.ndarray | None = None,
+    color: str | None = None,
+    linestyle: str | None = None,
+    linewidth: float | None = None,
 ) -> list[ContourItem]:
     """Create contour items
 
@@ -142,6 +145,11 @@ def create_contour_items(
          ``numpy.meshgrid``), or it must both be 1-D such that ``len(Y) == N``
          is the number of rows in *Z*.
          If none, they are assumed to be integer indices, i.e. ``Y = range(N)``.
+        color: contour color name. Default is None
+        linestyle: contour line style (MATLAB-like string or "SolidLine",
+         "DashLine", "DotLine", "DashDotLine", "DashDotDotLine", "NoPen").
+         Default is None
+        linewidth: contour line width (pixels). Default is None
 
     Returns:
         A list of :py:class:`.ContourItem` instances.
@@ -153,6 +161,14 @@ def create_contour_items(
         param = ShapeParam("Contour", icon="contour.png")
         item = ContourItem(points=cline.vertices, shapeparam=param)
         item.set_style("plot", "shape/contour")
+        if color is not None:
+            item.shapeparam.line.color = COLORS.get(color, color)
+        if linestyle is not None:
+            item.shapeparam.line.set_style_from_matlab(linestyle)
+        if linewidth is not None:
+            item.shapeparam.line.width = linewidth
+        if color is not None or linestyle is not None or linewidth is not None:
+            item.shapeparam.update_item(item)
         item.setTitle(_("Contour") + f"[Z={cline.level}]")
         items.append(item)
     return items

plotpy/tests/items/test_contours.py

@@ -28,7 +28,13 @@ def test_contours():
     """Contour plotting on two nearby 2D Gaussians"""
     with qt_app_context(exec_loop=True):
         image = create_two_gaussians()
-        contour_items = make.contours(image, np.linspace(0.2, 1.0, 5))
+        contour_items = make.contours(
+            image,
+            np.linspace(0.2, 1.0, 5),
+            color="y",
+            linestyle=":",
+            linewidth=2.0,
+        )
 
         assert contour_items
         assert all(isinstance(item, ContourItem) for item in contour_items)

plotpy/tests/unit/test_contour.py

@@ -1,7 +1,9 @@
 import unittest
 
 import numpy as np
+from qtpy.QtCore import Qt
 
+from plotpy.builder import make
 from plotpy.items.contour import compute_contours
 from plotpy.tests.data import gen_xyz_data
 
@@ -54,3 +56,26 @@ def test_contour_abs_level_1(self):
         for line in lines:
             assert np.all(np.equal(line.vertices[0], line.vertices[-1]))
             assert line.level == 1.0
+
+    def test_contour_item_style(self):
+        """Test contour builder line style parameters."""
+        items = make.contours(
+            self.Z,
+            1.0,
+            self.X,
+            self.Y,
+            color="r",
+            linestyle="--",
+            linewidth=2.5,
+        )
+        assert items
+        for item in items:
+            assert item.pen.color().name() == "#ff0000"
+            assert item.pen.style() == Qt.DashLine
+            assert item.pen.widthF() == 2.5
+
+        default_item = make.contours(self.Z, 1.0, self.X, self.Y)[0]
+        color_item = make.contours(self.Z, 1.0, self.X, self.Y, color="b")[0]
+        assert color_item.pen.color().name() == "#0000ff"
+        assert color_item.pen.style() == default_item.pen.style()
+        assert color_item.pen.widthF() == default_item.pen.widthF()