Fix: stackplot accepts CN colors without mutating Axes cycler (swev-id: matplotlib__matplotlib-24026)

lib/matplotlib/stackplot.py

@@ -6,12 +6,18 @@
 (https://stackoverflow.com/users/66549/doug)
 """
 
+import itertools
+import re
+
 import numpy as np
 
-from matplotlib import _api
+from matplotlib import _api, colors as mcolors
 
 __all__ = ['stackplot']
 
+_CN_ALIAS_RE = re.compile(r'C\d+')
+_SINGLE_LETTER_COLOR_CHARS = frozenset('bgrcmykwBGRCMYKW')
+
 
 def stackplot(axes, x, *args,
               labels=(), colors=None, baseline='zero',
@@ -69,8 +75,33 @@ def stackplot(axes, x, *args,
     y = np.row_stack(args)
 
     labels = iter(labels)
+    color_iter = None
     if colors is not None:
-        axes.set_prop_cycle(color=colors)
+        if isinstance(colors, str):
+            color_str = colors.strip()
+            if _CN_ALIAS_RE.fullmatch(color_str):
+                parsed_colors = [color_str]
+            elif ',' in color_str:
+                parsed_colors = [part.strip() for part in color_str.split(',')
+                                 if part.strip()]
+            elif (color_str and set(color_str)
+                  <= _SINGLE_LETTER_COLOR_CHARS):
+                parsed_colors = list(color_str)
+            else:
+                parsed_colors = [color_str]
+        else:
+            try:
+                iter(colors)
+            except TypeError:
+                parsed_colors = [colors]
+            else:
+                if isinstance(colors, (set, frozenset)):
+                    parsed_colors = [colors]
+                else:
+                    parsed_colors = list(colors)
+        rgba_list = mcolors.to_rgba_array(parsed_colors)
+        if len(rgba_list):
+            color_iter = itertools.cycle(rgba_list)
 
     # Assume data passed has not been 'stacked', so stack it here.
     # We'll need a float buffer for the upcoming calculations.
@@ -108,7 +139,9 @@ def stackplot(axes, x, *args,
         stack += first_line
 
     # Color between x = 0 and the first array.
-    color = axes._get_lines.get_next_color()
+    color = (next(color_iter)
+             if color_iter is not None
+             else axes._get_lines.get_next_color())
     coll = axes.fill_between(x, first_line, stack[0, :],
                              facecolor=color, label=next(labels, None),
                              **kwargs)
@@ -117,7 +150,9 @@ def stackplot(axes, x, *args,
 
     # Color between array i-1 and array i
     for i in range(len(y) - 1):
-        color = axes._get_lines.get_next_color()
+        color = (next(color_iter)
+                 if color_iter is not None
+                 else axes._get_lines.get_next_color())
         r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
                                    facecolor=color, label=next(labels, None),
                                    **kwargs))

lib/matplotlib/tests/test_stackplot_colors.py

@@ -0,0 +1,102 @@
+import warnings
+
+import numpy as np
+from cycler import cycler
+
+
+def _import_pyplot_and_colors():
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", DeprecationWarning)
+        from matplotlib import colors as mcolors, pyplot as plt
+    return mcolors, plt
+
+
+mcolors, plt = _import_pyplot_and_colors()
+
+
+def _stackplot_data(num_layers=3):
+    x = np.arange(3)
+    layers = [np.full_like(x, fill_value=i + 1, dtype=float)
+              for i in range(num_layers)]
+    return x, layers
+
+
+def _facecolors(collections):
+    return [tuple(coll.get_facecolors()[0]) for coll in collections]
+
+
+def test_stackplot_accepts_cn_colors():
+    fig, ax = plt.subplots()
+    x, layers = _stackplot_data()
+
+    collections = ax.stackplot(x, *layers, colors=['C2', 'C3', 'C4'])
+
+    actual = _facecolors(collections)
+    expected = [tuple(mcolors.to_rgba(color)) for color in ['C2', 'C3', 'C4']]
+    assert actual == expected
+
+
+def test_stackplot_does_not_mutate_axes_cycler():
+    fig, (ax_left, ax_right) = plt.subplots(1, 2)
+    x, layers = _stackplot_data(2)
+    prop_colors = ['tab:blue', 'tab:orange', 'tab:green']
+    ax_left.set_prop_cycle(color=prop_colors)
+    ax_right.set_prop_cycle(color=prop_colors)
+
+    ax_left.stackplot(x, *layers, colors=['C3', 'C4'])
+
+    left_line, = ax_left.plot(x, x)
+    right_line, = ax_right.plot(x, x)
+
+    assert left_line.get_color() == right_line.get_color() == prop_colors[0]
+
+
+def test_stackplot_colors_repeat_when_short():
+    fig, ax = plt.subplots()
+    x, layers = _stackplot_data(4)
+
+    collections = ax.stackplot(x, *layers, colors=['C1'])
+
+    actual = _facecolors(collections)
+    expected_color = tuple(mcolors.to_rgba('C1'))
+    assert actual == [expected_color] * len(layers)
+
+
+def test_stackplot_colors_exact_match():
+    fig, ax = plt.subplots()
+    x, layers = _stackplot_data(3)
+    requested = ['C5', 'C6', 'C7']
+
+    collections = ax.stackplot(x, *layers, colors=requested)
+
+    actual = _facecolors(collections)
+    expected = [tuple(mcolors.to_rgba(color)) for color in requested]
+    assert actual == expected
+
+
+def test_stackplot_colors_stringlist_single_letters():
+    fig, ax = plt.subplots()
+    x, layers = _stackplot_data(3)
+
+    collections = ax.stackplot(x, *layers, colors='rgb')
+
+    actual = _facecolors(collections)
+    expected = [tuple(mcolors.to_rgba(color)) for color in ['r', 'g', 'b']]
+    assert actual == expected
+
+
+def test_stackplot_uses_axes_cycler_when_colors_none():
+    fig, ax = plt.subplots()
+    x, layers = _stackplot_data(3)
+    cycler_colors = ['tab:blue', 'tab:orange', 'tab:green', 'tab:red']
+    ax.set_prop_cycle(cycler(color=cycler_colors))
+
+    collections = ax.stackplot(x, *layers)
+
+    actual = _facecolors(collections)
+    expected = [tuple(mcolors.to_rgba(color))
+                for color in cycler_colors[:len(layers)]]
+    assert actual == expected
+
+    next_line, = ax.plot(x, x)
+    assert next_line.get_color() == cycler_colors[len(layers)]