Support pandas.DataFrame as input

timhoffm · Jan 18, 2025 · df4c389 · df4c389
1 parent 2e0660a
commit df4c389
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 4 deletions.
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -3148,6 +3148,20 @@ def grouped_bar(self, heights, *, positions=None, group_spacing=1.5, bar_spacing
                   df = pd.DataFrame(array, index=group_labels, columns=dataset_labels)
                   df.plot.bar()
 
+            - a `pandas.DataFrame`.
+
+              .. code-block::
+
+                  df = pd.DataFrame(
+                      np.random.random((2, 3))
+                      index=["group_A", "group_B"],
+                      columns=["dataset_0", "dataset_1", "dataset_2"]
+                  )
+                  grouped_bar(df)
+
+              Note that ``grouped_bar(df)`` produced a structurally equivalent plot like
+              ``df.plot.bar()`.
+
         positions : array-like, optional
             The center positions of the bar groups. The values have to be equidistant.
             If not given, a sequence of integer positions 0, 1, 2, ... is used.
@@ -3198,13 +3212,19 @@ def grouped_bar(self, heights, *, positions=None, group_spacing=1.5, bar_spacing
                   See also `.Artist.remove()`.
 
         """
-        if hasattr(heights, 'keys'):
+        if cbook._is_pandas_dataframe(heights):
+            if labels is None:
+                labels = heights.columns.tolist()
+            if tick_labels is None:
+                tick_labels = heights.index.tolist()
+            heights = heights.to_numpy().T
+        elif hasattr(heights, 'keys'):  # dict
             if labels is not None:
                 raise ValueError(
                     "'labels' cannot be used if 'heights' are a mapping")
             labels = heights.keys()
             heights = list(heights.values())
-        elif hasattr(heights, 'shape'):
+        elif hasattr(heights, 'shape'):  # numpy array
             heights = heights.T
 
         num_datasets = len(heights)

diff --git a/lib/matplotlib/axes/_axes.pyi b/lib/matplotlib/axes/_axes.pyi
@@ -287,7 +287,7 @@ class Axes(_AxesBase):
     ) -> PolyCollection: ...
     def grouped_bar(
         self,
-        heights : Sequence[ArrayLike] | dict[str, ArrayLike] | np.ndarray,
+        heights : Sequence[ArrayLike] | dict[str, ArrayLike] | np.ndarray | pd.DataFrame,
         *,
         positions : ArrayLike | None = ...,
         tick_labels : Sequence[str] | None = ...,

diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -94,6 +94,7 @@
 
     import PIL.Image
     from numpy.typing import ArrayLike
+    import pandas as pd
 
     import matplotlib.axes
     import matplotlib.artist
@@ -3391,7 +3392,7 @@ def grid(
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @_copy_docstring_and_deprecators(Axes.grouped_bar)
 def grouped_bar(
-    heights: Sequence[ArrayLike] | dict[str, ArrayLike] | np.ndarray,
+    heights: Sequence[ArrayLike] | dict[str, ArrayLike] | np.ndarray | pd.DataFrame,
     *,
     positions: ArrayLike | None = None,
     group_spacing: float | None = 1.5,