Feature: Support passing DataFrames to table.table (#28830)

anijjar · timhoffm · web-flow · commit 8e6d6b6bb27b · 2024-10-29T11:01:02.000+01:00
---------

Co-authored-by: Tim Hoffmann &lt;2836374+timhoffm@users.noreply.github.com&gt;
diff --git a/doc/users/next_whats_new/pass_pandasDataFrame_into_table.rst b/doc/users/next_whats_new/pass_pandasDataFrame_into_table.rst
@@ -0,0 +1,20 @@
+``ax.table`` will accept a pandas DataFrame
+--------------------------------------------
+
+The `~.axes.Axes.table` method can now accept a Pandas DataFrame for the ``cellText`` argument.
+
+.. code-block:: python
+
+    import matplotlib.pyplot as plt
+    import pandas as pd
+
+    data = {
+        'Letter': ['A', 'B', 'C'],
+        'Number': [100, 200, 300]
+    }
+
+    df = pd.DataFrame(data)
+    fig, ax = plt.subplots()
+    table = ax.table(df, loc='center')  # or table = ax.table(cellText=df, loc='center')
+    ax.axis('off')
+    plt.show()
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -2391,3 +2391,15 @@ def _auto_format_str(fmt, value):
         return fmt % (value,)
     except (TypeError, ValueError):
         return fmt.format(value)
+
+
+def _is_pandas_dataframe(x):
+    """Check if 'x' is a Pandas DataFrame."""
+    try:
+        # we're intentionally not attempting to import Pandas. If somebody
+        # has created a Pandas DataFrame, Pandas should already be in sys.modules
+        return isinstance(x, sys.modules['pandas'].DataFrame)
+    except Exception:  # TypeError, KeyError, AttributeError, maybe others?
+        # we're attempting to access attributes on imported modules which
+        # may have arbitrary user code, so we deliberately catch all exceptions
+        return False
diff --git a/lib/matplotlib/table.py b/lib/matplotlib/table.py
@@ -33,6 +33,8 @@
 from .transforms import Bbox
 from .path import Path
 
+from .cbook import _is_pandas_dataframe
+
 
 class Cell(Rectangle):
     """
@@ -670,7 +672,7 @@ def table(ax,
 
     Parameters
     ----------
-    cellText : 2D list of str, optional
+    cellText : 2D list of str or pandas.DataFrame, optional
         The texts to place into the table cells.
 
         *Note*: Line breaks in the strings are currently not accounted for and
@@ -740,6 +742,21 @@ def table(ax,
         cols = len(cellColours[0])
         cellText = [[''] * cols] * rows
 
+    # Check if we have a Pandas DataFrame
+    if _is_pandas_dataframe(cellText):
+        # if rowLabels/colLabels are empty, use DataFrame entries.
+        # Otherwise, throw an error.
+        if rowLabels is None:
+            rowLabels = cellText.index
+        else:
+            raise ValueError("rowLabels cannot be used alongside Pandas DataFrame")
+        if colLabels is None:
+            colLabels = cellText.columns
+        else:
+            raise ValueError("colLabels cannot be used alongside Pandas DataFrame")
+        # Update cellText with only values
+        cellText = cellText.values
+
     rows = len(cellText)
     cols = len(cellText[0])
     for row in cellText:
diff --git a/lib/matplotlib/table.pyi b/lib/matplotlib/table.pyi
@@ -8,7 +8,9 @@ from .transforms import Bbox
 from .typing import ColorType
 
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import Any, Literal, TYPE_CHECKING
+
+from pandas import DataFrame
 
 class Cell(Rectangle):
     PAD: float
@@ -68,7 +70,7 @@ class Table(Artist):
 
 def table(
     ax: Axes,
-    cellText: Sequence[Sequence[str]] | None = ...,
+    cellText: Sequence[Sequence[str]] | DataFrame | None = ...,
     cellColours: Sequence[Sequence[ColorType]] | None = ...,
     cellLoc: Literal["left", "center", "right"] = ...,
     colWidths: Sequence[float] | None = ...,
diff --git a/lib/matplotlib/tests/test_table.py b/lib/matplotlib/tests/test_table.py
@@ -253,3 +253,20 @@ def __repr__(self):
 
     munits.registry.pop(FakeUnit)
     assert not munits.registry.get_converter(FakeUnit)
+
+
+def test_table_dataframe(pd):
+    # Test if Pandas Data Frame can be passed in cellText
+
+    data = {
+        'Letter': ['A', 'B', 'C'],
+        'Number': [100, 200, 300]
+    }
+
+    df = pd.DataFrame(data)
+    fig, ax = plt.subplots()
+    table = ax.table(df, loc='center')
+
+    for r, (index, row) in enumerate(df.iterrows()):
+        for c, col in enumerate(df.columns if r == 0 else row.values):
+            assert table[r if r == 0 else r+1, c].get_text().get_text() == str(col)
