Merge pull request matplotlib#30808 from rcomer/style-iterator

Consolidate style parameter handling for plotting methods that call other plotting methods

Author: tacaswell

doc/release/next_whats_new/stackplot_style_sequences.rst

@@ -0,0 +1,6 @@
+Stackplot styling
+-----------------
+
+`~.Axes.stackplot` now accepts sequences for the style parameters *facecolor*,
+*edgecolor*, *linestyle*, and *linewidth*, similar to how the *hatch* parameter
+is already handled.

lib/matplotlib/_style_helpers.py

@@ -0,0 +1,51 @@
+import collections.abc
+import itertools
+
+import numpy as np
+
+import matplotlib.cbook as cbook
+import matplotlib.colors as mcolors
+import matplotlib.lines as mlines
+
+
+def check_non_empty(key, value):
+    """Raise a TypeError if an empty sequence is passed"""
+    if (not cbook.is_scalar_or_string(value) and
+            isinstance(value, collections.abc.Sized) and len(value) == 0):
+        raise TypeError(f'{key} must not be an empty sequence')
+
+
+def style_generator(kw):
+    """
+    Helper for handling style sequences (e.g. facecolor=['r', 'b', 'k']) within plotting
+    methods that repeatedly call other plotting methods (e.g. hist, stackplot).  Remove
+    style keywords from the given dictionary.  Return the reduced dictionary together
+    with a generator which provides a series of dictionaries to be used in each call to
+    the wrapped function.
+    """
+    kw_iterators = {}
+    remaining_kw = {}
+    for key, value in kw.items():
+        if key in ['facecolor', 'edgecolor']:
+            if value is None or cbook._str_lower_equal(value, 'none'):
+                kw_iterators[key] = itertools.repeat(value)
+            else:
+                check_non_empty(key, value)
+                kw_iterators[key] = itertools.cycle(mcolors.to_rgba_array(value))
+
+        elif key in ['hatch', 'linewidth']:
+            check_non_empty(key, value)
+            kw_iterators[key] = itertools.cycle(np.atleast_1d(value))
+
+        elif key == 'linestyle':
+            check_non_empty(key, value)
+            kw_iterators[key] = itertools.cycle(mlines._get_dash_patterns(value))
+
+        else:
+            remaining_kw[key] = value
+
+    def style_gen():
+        while True:
+            yield {key: next(val) for key, val in kw_iterators.items()}
+
+    return remaining_kw, style_gen()

lib/matplotlib/axes/_axes.py

@@ -32,7 +32,7 @@
 import matplotlib.transforms as mtransforms
 import matplotlib.tri as mtri
 import matplotlib.units as munits
-from matplotlib import _api, _docstring, _preprocess_data
+from matplotlib import _api, _docstring, _preprocess_data, _style_helpers
 from matplotlib.axes._base import (
     _AxesBase, _TransformedBoundsLocator, _process_plot_format)
 from matplotlib.axes._secondary_axes import SecondaryAxis
@@ -3194,6 +3194,16 @@ def grouped_bar(self, heights, *, positions=None, group_spacing=1.5, bar_spacing
 
         **kwargs : `.Rectangle` properties
 
+            Properties applied to all bars. The following properties additionally
+            accept a sequence of values corresponding to the datasets in
+            *heights*:
+
+            - *edgecolor*
+            - *facecolor*
+            - *linewidth*
+            - *linestyle*
+            - *hatch*
+
             %(Rectangle:kwdoc)s
 
         Returns
@@ -3320,6 +3330,8 @@ def grouped_bar(self, heights, *, positions=None, group_spacing=1.5, bar_spacing
             # TODO: do we want to be more restrictive and check lengths?
             colors = itertools.cycle(colors)
 
+        kwargs, style_gen = _style_helpers.style_generator(kwargs)
+
         bar_width = (group_distance /
                      (num_datasets + (num_datasets - 1) * bar_spacing + group_spacing))
         bar_spacing_abs = bar_spacing * bar_width
@@ -3333,15 +3345,16 @@ def grouped_bar(self, heights, *, positions=None, group_spacing=1.5, bar_spacing
         # place the bars, but only use numerical positions, categorical tick labels
         # are handled separately below
         bar_containers = []
-        for i, (hs, label, color) in enumerate(zip(heights, labels, colors)):
+        for i, (hs, label, color, styles) in enumerate(zip(heights, labels, colors,
+                                                           style_gen)):
             lefts = (group_centers - 0.5 * group_distance + margin_abs
                      + i * (bar_width + bar_spacing_abs))
             if orientation == "vertical":
                 bc = self.bar(lefts, hs, width=bar_width, align="edge",
-                              label=label, color=color, **kwargs)
+                              label=label, color=color, **styles, **kwargs)
             else:
                 bc = self.barh(lefts, hs, height=bar_width, align="edge",
-                               label=label, color=color, **kwargs)
+                               label=label, color=color, **styles, **kwargs)
             bar_containers.append(bc)
 
         if tick_labels is not None:
@@ -7632,38 +7645,15 @@ def hist(self, x, bins=None, range=None, density=False, weights=None,
         labels = [] if label is None else np.atleast_1d(np.asarray(label, str))
 
         if histtype == "step":
-            ec = kwargs.get('edgecolor', colors)
-        else:
-            ec = kwargs.get('edgecolor', None)
-        if ec is None or cbook._str_lower_equal(ec, 'none'):
-            edgecolors = itertools.repeat(ec)
-        else:
-            edgecolors = itertools.cycle(mcolors.to_rgba_array(ec))
-
-        fc = kwargs.get('facecolor', colors)
-        if cbook._str_lower_equal(fc, 'none'):
-            facecolors = itertools.repeat(fc)
-        else:
-            facecolors = itertools.cycle(mcolors.to_rgba_array(fc))
+            kwargs.setdefault('edgecolor', colors)
 
-        hatches = itertools.cycle(np.atleast_1d(kwargs.get('hatch', None)))
-        linewidths = itertools.cycle(np.atleast_1d(kwargs.get('linewidth', None)))
-        if 'linestyle' in kwargs:
-            linestyles = itertools.cycle(mlines._get_dash_patterns(kwargs['linestyle']))
-        else:
-            linestyles = itertools.repeat(None)
+        kwargs, style_gen = _style_helpers.style_generator(kwargs)
 
         for patch, lbl in itertools.zip_longest(patches, labels):
             if not patch:
                 continue
             p = patch[0]
-            kwargs.update({
-                'hatch': next(hatches),
-                'linewidth': next(linewidths),
-                'linestyle': next(linestyles),
-                'edgecolor': next(edgecolors),
-                'facecolor': next(facecolors),
-            })
+            kwargs.update(next(style_gen))
             p._internal_update(kwargs)
             if lbl is not None:
                 p.set_label(lbl)

lib/matplotlib/meson.build

@@ -17,6 +17,7 @@ python_sources = [
   '_mathtext.py',
   '_mathtext_data.py',
   '_pylab_helpers.py',
+  '_style_helpers.py',
   '_text_helpers.py',
   '_tight_bbox.py',
   '_tight_layout.py',

lib/matplotlib/pyplot.py

@@ -4208,15 +4208,12 @@ def spy(
 
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @_copy_docstring_and_deprecators(Axes.stackplot)
-def stackplot(
-    x, *args, labels=(), colors=None, hatch=None, baseline="zero", data=None, **kwargs
-):
+def stackplot(x, *args, labels=(), colors=None, baseline="zero", data=None, **kwargs):
     return gca().stackplot(
         x,
         *args,
         labels=labels,
         colors=colors,
-        hatch=hatch,
         baseline=baseline,
         **({"data": data} if data is not None else {}),
         **kwargs,

lib/matplotlib/stackplot.py

@@ -6,17 +6,16 @@
 (https://stackoverflow.com/users/66549/doug)
 """
 
-import itertools
 
 import numpy as np
 
-from matplotlib import _api
+from matplotlib import cbook, collections, _api, _style_helpers
 
 __all__ = ['stackplot']
 
 
 def stackplot(axes, x, *args,
-              labels=(), colors=None, hatch=None, baseline='zero',
+              labels=(), colors=None, baseline='zero',
               **kwargs):
     """
     Draw a stacked area plot or a streamgraph.
@@ -55,23 +54,26 @@ def stackplot(axes, x, *args,
 
         If not specified, the colors from the Axes property cycle will be used.
 
-    hatch : list of str, default: None
-        A sequence of hatching styles.  See
-        :doc:`/gallery/shapes_and_collections/hatch_style_reference`.
-        The sequence will be cycled through for filling the
-        stacked areas from bottom to top.
-        It need not be exactly the same length as the number
-        of provided *y*, in which case the styles will repeat from the
-        beginning.
-
-        .. versionadded:: 3.9
-           Support for list input
-
     data : indexable object, optional
         DATA_PARAMETER_PLACEHOLDER
 
     **kwargs
-        All other keyword arguments are passed to `.Axes.fill_between`.
+        All other keyword arguments are passed to `.Axes.fill_between`.  The
+            following parameters additionally accept a sequence of values
+            corresponding to the *y* datasets:
+
+            - *hatch*
+            - *edgecolor*
+            - *facecolor*
+            - *linewidth*
+            - *linestyle*
+
+            .. versionadded:: 3.9
+               Allowing a sequence of strings for *hatch*.
+
+            .. versionadded:: 3.11
+               Allowing sequences of values in above listed `.Axes.fill_between`
+               parameters.
 
     Returns
     -------
@@ -83,15 +85,13 @@ def stackplot(axes, x, *args,
     y = np.vstack(args)
 
     labels = iter(labels)
-    if colors is not None:
-        colors = itertools.cycle(colors)
-    else:
-        colors = (axes._get_lines.get_next_color() for _ in y)
+    if colors is None:
+        colors = [axes._get_lines.get_next_color() for _ in y]
+
+    kwargs = cbook.normalize_kwargs(kwargs, collections.PolyCollection)
+    kwargs.setdefault('facecolor', colors)
 
-    if hatch is None or isinstance(hatch, str):
-        hatch = itertools.cycle([hatch])
-    else:
-        hatch = itertools.cycle(hatch)
+    kwargs, style_gen = _style_helpers.style_generator(kwargs)
 
     # Assume data passed has not been 'stacked', so stack it here.
     # We'll need a float buffer for the upcoming calculations.
@@ -130,18 +130,14 @@ def stackplot(axes, x, *args,
 
     # Color between x = 0 and the first array.
     coll = axes.fill_between(x, first_line, stack[0, :],
-                             facecolor=next(colors),
-                             hatch=next(hatch),
                              label=next(labels, None),
-                             **kwargs)
+                             **next(style_gen), **kwargs)
     coll.sticky_edges.y[:] = [0]
     r = [coll]
 
     # Color between array i-1 and array i
     for i in range(len(y) - 1):
         r.append(axes.fill_between(x, stack[i, :], stack[i + 1, :],
-                                   facecolor=next(colors),
-                                   hatch=next(hatch),
                                    label=next(labels, None),
-                                   **kwargs))
+                                   **next(style_gen), **kwargs))
     return r

lib/matplotlib/tests/test__style_helpers.py

@@ -0,0 +1,83 @@
+import pytest
+
+import matplotlib.colors as mcolors
+from matplotlib.lines import _get_dash_pattern
+from matplotlib._style_helpers import style_generator
+
+
+@pytest.mark.parametrize('key, value', [('facecolor', ["b", "g", "r"]),
+                                        ('edgecolor', ["b", "g", "r"]),
+                                        ('hatch', ["/", "\\", "."]),
+                                        ('linestyle', ["-", "--", ":"]),
+                                        ('linewidth', [1, 1.5, 2])])
+def test_style_generator_list(key, value):
+    """Test that style parameter lists are distributed to the generator."""
+    kw = {'foo': 12, key: value}
+    new_kw, gen = style_generator(kw)
+
+    assert new_kw == {'foo': 12}
+
+    for v in value * 2:  # Result should repeat
+        style_dict = next(gen)
+        assert len(style_dict) == 1
+        if key.endswith('color'):
+            assert mcolors.same_color(v, style_dict[key])
+        elif key == 'linestyle':
+            assert _get_dash_pattern(v) == style_dict[key]
+        else:
+            assert v == style_dict[key]
+
+
+@pytest.mark.parametrize('key, value', [('facecolor', "b"),
+                                        ('edgecolor', "b"),
+                                        ('hatch', "/"),
+                                        ('linestyle', "-"),
+                                        ('linewidth', 1)])
+def test_style_generator_single(key, value):
+    """Test that single-value style parameters are distributed to the generator."""
+    kw = {'foo': 12, key: value}
+    new_kw, gen = style_generator(kw)
+
+    assert new_kw == {'foo': 12}
+    for _ in range(2):  # Result should repeat
+        style_dict = next(gen)
+        if key.endswith('color'):
+            assert mcolors.same_color(value, style_dict[key])
+        elif key == 'linestyle':
+            assert _get_dash_pattern(value) == style_dict[key]
+        else:
+            assert value == style_dict[key]
+
+
+@pytest.mark.parametrize('key', ['facecolor', 'hatch', 'linestyle'])
+def test_style_generator_raises_on_empty_style_parameter_list(key):
+    kw = {key: []}
+    with pytest.raises(TypeError, match=f'{key} must not be an empty sequence'):
+        style_generator(kw)
+
+
+def test_style_generator_sequence_type_styles():
+    """
+    Test that sequence type style values are detected as single value
+    and passed to a all elements of the generator.
+    """
+    kw = {'facecolor':  ('r', 0.5),
+          'edgecolor': [0.5, 0.5, 0.5],
+          'linestyle': (0, (1, 1))}
+
+    _, gen = style_generator(kw)
+    for _ in range(2):  # Result should repeat
+        style_dict = next(gen)
+        mcolors.same_color(kw['facecolor'], style_dict['facecolor'])
+        mcolors.same_color(kw['edgecolor'], style_dict['edgecolor'])
+        kw['linestyle'] == style_dict['linestyle']
+
+
+def test_style_generator_none():
+    kw = {'facecolor': 'none',
+          'edgecolor': 'none'}
+    _, gen = style_generator(kw)
+    for _ in range(2):  # Result should repeat
+        style_dict = next(gen)
+        assert style_dict['facecolor'] == 'none'
+        assert style_dict['edgecolor'] == 'none'