Add colors and kwargs, typing stubs, pyplot wrapper

timhoffm · timhoffm · commit 128b50957288 · 2024-09-11T13:45:41.000+02:00
diff --git a/galleries/examples/lines_bars_and_markers/grouped_bar_chart.py b/galleries/examples/lines_bars_and_markers/grouped_bar_chart.py
@@ -145,6 +145,25 @@
 axs[1, 1].grouped_bar(x, data, group_spacing=0.5, bar_spacing=0.1)
 
 
+# %%
+# Styling
+# -------
+# The bars can be styled through additional keyword arguments. Currently,
+# the only per-dataset setting is ``colors``. Additionally, all
+# `.Rectangle parameters` are passed through and applied to all datasets.
+
+x = ['A', 'B', 'C']
+data = {
+    'data1': [1, 2, 3],
+    'data2': [1.2, 2.2, 3.2],
+    'data3': [1.4, 2.4, 3.4],
+    'data4': [1.6, 2.6, 3.6],
+}
+
+fig, ax = plt.subplots()
+ax.grouped_bar(x, data, colors=["r", "g", "b", "m"], edgecolor="black")
+
+
 # %%
 # Horizontal grouped bars
 # -----------------------
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3014,10 +3014,17 @@ def broken_barh(self, xranges, yrange, **kwargs):
         return col
 
     def grouped_bar(self, x, heights, *, group_spacing=1.5, bar_spacing=0,
-                    dataset_labels=None, orientation="vertical"):
+                    dataset_labels=None, orientation="vertical", colors=None,
+                    **kwargs):
         """
+        Make a grouped bar plot.
+
+        .. note::
+            This function is new in v3.10, and the API is still provisional.
+            We may still fine-tune some aspects based on user-feedback.
+
         Parameters
-        -----------
+        ----------
         x : array-like or list of str
             The center positions of the bar groups. If these are numeric values,
             they have to be equidistant. As with `~.Axes.bar`, you can provide
@@ -3081,9 +3088,20 @@ def grouped_bar(self, x, heights, *, group_spacing=1.5, bar_spacing=0,
             The labels of the datasets.
 
         orientation : {"vertical", "horizontal"}, default: vertical
-        """
-        _api.check_in_list(["vertical", "horizontal"], orientation=orientation)
 
+        colors : list of :mpltype:`color`, optional
+            A sequence of colors to be cycled through and used to color bars
+            of the different datasets. The sequence need not be exactly the
+            same length as the number of provided y, in which case the colors
+            will repeat from the beginning.
+
+            If not specified, the colors from the Axes property cycle will be used.
+
+        **kwargs : `.Rectangle` properties
+
+            %(Rectangle:kwdoc)s
+
+        """
         if hasattr(heights, 'keys'):
             if dataset_labels is not None:
                 raise ValueError(
@@ -3118,6 +3136,15 @@ def grouped_bar(self, x, heights, *, group_spacing=1.5, bar_spacing=0,
                     f"has {len(dataset)} groups"
                 )
 
+        _api.check_in_list(["vertical", "horizontal"], orientation=orientation)
+
+        if colors is None:
+            colors = itertools.cycle([None])
+        else:
+            # Note: This is equivalent to the behavior in stackplot
+            # TODO: do we want to be more restrictive and check lengths?
+            colors = itertools.cycle(colors)
+
         bar_width = (group_distance /
                      (num_datasets + (num_datasets - 1) * bar_spacing + group_spacing))
         bar_spacing_abs = bar_spacing * bar_width
@@ -3130,15 +3157,16 @@ def grouped_bar(self, x, heights, *, group_spacing=1.5, bar_spacing=0,
 
         # place the bars, but only use numerical positions, categorical tick labels
         # are handled separately below
-        for i, (hs, dataset_label) in enumerate(zip(heights, dataset_labels)):
+        for i, (hs, dataset_label, color) in enumerate(
+                zip(heights, dataset_labels, colors)):
             lefts = (group_centers - 0.5 * group_distance + margin_abs
                      + i * (bar_width + bar_spacing_abs))
             if orientation == "vertical":
                 self.bar(lefts, hs, width=bar_width, align="edge",
-                         label=dataset_label)
+                         label=dataset_label, color=color, **kwargs)
             else:
                 self.barh(lefts, hs, height=bar_width, align="edge",
-                          label=dataset_label)
+                          label=dataset_label, color=color, **kwargs)
 
         if tick_labels is not None:
             if orientation == "vertical":
diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -284,6 +284,18 @@ class Axes(_AxesBase):
         data=...,
         **kwargs
     ) -> PolyCollection: ...
+    def grouped_bar(
+        self,
+        x : ArrayLike,
+        heights : ArrayLike,
+        *,
+        group_spacing : float | None = ...,
+        bar_spacing : float | None = ...,
+        dataset_labels : Sequence[str] | None = ...,
+        orientation: Literal["vertical", "horizontal"] = ...,
+        colors: Iterable[ColorType] | None = ...,
+        **kwargs
+    ) -> None: ...
     def stem(
         self,
         *args: ArrayLike | str,
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -3376,6 +3376,31 @@ def grid(
     gca().grid(visible=visible, which=which, axis=axis, **kwargs)
 
 
+# Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
+@_copy_docstring_and_deprecators(Axes.grouped_bar)
+def grouped_bar(
+    x: ArrayLike,
+    heights: ArrayLike,
+    *,
+    group_spacing: float | None = 1.5,
+    bar_spacing: float | None = 0,
+    dataset_labels: Sequence[str] | None = None,
+    orientation: Literal["vertical", "horizontal"] = "vertical",
+    colors: Iterable[ColorType] | None = None,
+    **kwargs,
+) -> None:
+    gca().grouped_bar(
+        x,
+        heights,
+        group_spacing=group_spacing,
+        bar_spacing=bar_spacing,
+        dataset_labels=dataset_labels,
+        orientation=orientation,
+        colors=colors,
+        **kwargs,
+    )
+
+
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @_copy_docstring_and_deprecators(Axes.hexbin)
 def hexbin(
diff --git a/tools/boilerplate.py b/tools/boilerplate.py
@@ -238,6 +238,7 @@ def boilerplate_gen():
         'fill_between',
         'fill_betweenx',
         'grid',
+        'grouped_bar',
         'hexbin',
         'hist',
         'stairs',
