Handle single colors in ContourSet

rcomer · rcomer · commit 5ff1f17c4fe2 · 2024-09-08T10:58:30.000+01:00
diff --git a/doc/users/next_whats_new/contour_color.rst b/doc/users/next_whats_new/contour_color.rst
@@ -0,0 +1,21 @@
+Specifying a single color in ``contour`` and ``contourf``
+---------------------------------------------------------
+
+`~.Axes.contour` and `~.Axes.contourf` previously accepted a single color
+provided it was expressed as a string.  This restriction has now been removed
+and a single color in any format described in the :ref:`colors_def` tutorial
+may be passed.
+
+.. plot::
+    :include-source: true
+    :alt: Two-panel example contour plots.  The left panel has all transparent red contours.  The right panel has all dark blue contours.
+
+    import matplotlib.pyplot as plt
+
+    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(6, 3))
+    z = [[0, 1], [1, 2]]
+
+    ax1.contour(z, colors=('r', 0.4))
+    ax2.contour(z, colors=(0.1, 0.2, 0.5))
+
+    plt.show()
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -702,6 +702,11 @@ def __init__(self, ax, *args,
         self._extend_min = self.extend in ['min', 'both']
         self._extend_max = self.extend in ['max', 'both']
         if self.colors is not None:
+            if mcolors.is_color_like(self.colors):
+                color_sequence = [self.colors]
+            else:
+                color_sequence = self.colors
+
             ncolors = len(self.levels)
             if self.filled:
                 ncolors -= 1
@@ -718,19 +723,19 @@ def __init__(self, ax, *args,
             total_levels = (ncolors +
                             int(self._extend_min) +
                             int(self._extend_max))
-            if (len(self.colors) == total_levels and
+            if (len(color_sequence) == total_levels and
                     (self._extend_min or self._extend_max)):
                 use_set_under_over = True
                 if self._extend_min:
                     i0 = 1
 
-            cmap = mcolors.ListedColormap(self.colors[i0:None], N=ncolors)
+            cmap = mcolors.ListedColormap(color_sequence[i0:None], N=ncolors)
 
             if use_set_under_over:
                 if self._extend_min:
-                    cmap.set_under(self.colors[0])
+                    cmap.set_under(color_sequence[0])
                 if self._extend_max:
-                    cmap.set_over(self.colors[-1])
+                    cmap.set_over(color_sequence[-1])
 
         # label lists must be initialized here
         self.labelTexts = []
@@ -1498,10 +1503,12 @@ def _initialize_x_y(self, z):
     The sequence is cycled for the levels in ascending order. If the
     sequence is shorter than the number of levels, it's repeated.
 
-    As a shortcut, single color strings may be used in place of
-    one-element lists, i.e. ``'red'`` instead of ``['red']`` to color
-    all levels with the same color. This shortcut does only work for
-    color strings, not for other ways of specifying colors.
+    As a shortcut, single color may be used in place of one-element lists, i.e.
+    ``'red'`` instead of ``['red']`` to color all levels with the same color.
+
+    .. versionchanged:: 3.10
+        Previously a single color had to be expressed as a string, but now any
+        valid color format may be passed.
 
     By default (value *None*), the colormap specified by *cmap*
     will be used.
@@ -1569,10 +1576,10 @@ def _initialize_x_y(self, z):
 
         An existing `.QuadContourSet` does not get notified if
         properties of its colormap are changed. Therefore, an explicit
-        call `.QuadContourSet.changed()` is needed after modifying the
+        call ``QuadContourSet.changed()`` is needed after modifying the
         colormap. The explicit call can be left out, if a colorbar is
         assigned to the `.QuadContourSet` because it internally calls
-        `.QuadContourSet.changed()`.
+        ``QuadContourSet.changed()``.
 
     Example::
 
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -171,6 +171,18 @@ def test_given_colors_levels_and_extends():
         plt.colorbar(c, ax=ax)
 
 
+@pytest.mark.parametrize('color', ['green', ('r', 0.5), (0.1, 0.2, 0.5, 0.3)])
+@pytest.mark.parametrize('extend', ['neither', 'both'])
+@pytest.mark.parametrize('levels', [[0.5, 1, 1.5], [0.5, 0.75, 1, 1.25, 1.5]])
+def test_single_color_and_extend(color, levels, extend):
+    z = [[0, 1], [1, 2]]
+
+    _, ax = plt.subplots()
+    cs = ax.contour(z, levels=levels, colors=color, extend=extend)
+    for c in cs.get_edgecolors():
+        assert same_color(c, color)
+
+
 @image_comparison(['contour_log_locator.svg'], style='mpl20', remove_text=False)
 def test_log_locator_levels():
 
