[ENH] Vectorized Wedgeprops argument to pie

story645 · story645 · commit acfcd48e8181 · 2022-11-15T16:51:57.000-05:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -1,3 +1,4 @@
+import collections
 import functools
 import itertools
 import logging
@@ -3116,11 +3117,12 @@ def pie(self, x, explode=None, labels=None, colors=None,
         counterclock : bool, default: True
             Specify fractions direction, clockwise or counterclockwise.
 
-        wedgeprops : dict, default: None
-            Dict of arguments passed to the wedge objects making the pie.
-            For example, you can pass in ``wedgeprops = {'linewidth': 3}``
-            to set the width of the wedge border lines equal to 3.
-            For more details, look at the doc/arguments of the wedge object.
+        wedgeprops : dict, list of dict, default: None
+            Dict of arguments passed to each `matplotlib.patches.Wedge` of the
+            pie. For example, you can pass in ``wedgeprops = {'linewidth': 3}``
+            to set the width of the wedge border lines equal to 3. You can also
+            pass in a list of dicts to set each individual wedge object.
+            For a list of properties, see `matplotlib.patches.Wedge`.
             By default ``clip_on=False``.
 
         textprops : dict, default: None
@@ -3205,7 +3207,14 @@ def get_next_color():
         slices = []
         autotexts = []
 
-        for frac, label, expl in zip(x, labels, explode):
+        if isinstance(wedgeprops, collections.abc.Sequence):
+            if len(wedgeprops) != len(x):
+                raise ValueError(f'wedgeprops length={len(wedgeprops)}'
+                                 f' and input length={len(x)} are not equal')
+        else:
+            wedgeprops = itertools.repeat(wedgeprops)
+
+        for frac, label, expl, wprops in zip(x, labels, explode, wedgeprops):
             x, y = center
             theta2 = (theta1 + frac) if counterclock else (theta1 - frac)
             thetam = 2 * np.pi * 0.5 * (theta1 + theta2)
@@ -3217,7 +3226,7 @@ def get_next_color():
                                facecolor=get_next_color(),
                                clip_on=False,
                                label=label)
-            w.set(**wedgeprops)
+            w.set(**wprops)
             slices.append(w)
             self.add_patch(w)
 
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -5709,6 +5709,24 @@ def test_normalize_kwarg_pie():
     assert abs(t2[0][-1].theta2 - 360.) > 1e-3
 
 
+def test_pie_wedgeprops_error():
+    fig, ax = plt.subplots()
+    x = [0.3, 0.3, 0.1]
+    wedgeprops = [{'a': 1}]
+    with pytest.raises(ValueError, match=r"length=1 .* length=3"):
+        ax.pie(x, wedgeprops=wedgeprops)
+
+
+@check_figures_equal()
+def test_pie_multi_wedgeprops(fig_test, fig_ref):
+    x = [0.3, 0.3, 0.1]
+    wedgeprops = [{'hatch': '/'}, {'hatch': '+'}, {'hatch': '.'}]
+    fig_test.subplots().pie(x, wedgeprops=wedgeprops)
+    wedges, _ = fig_ref.subplots().pie(x)
+    for w, wp in zip(wedges, wedgeprops):
+        w.set(**wp)
+
+
 @image_comparison(['set_get_ticklabels.png'])
 def test_set_get_ticklabels():
     # test issue 2246
