FIX: make reshape2d accept pandas df with string indices

jklymak · jklymak · commit cf9bb88d0134 · 2020-08-29T15:12:26.000-07:00
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1355,6 +1355,17 @@ def _reshape_2D(X, name):
 
     *name* is used to generate the error message for invalid inputs.
     """
+
+    # unpack if we have a values or to_numpy method.
+    try:
+        X = X.to_numpy()
+    except AttributeError:
+        try:
+            if isinstance(X.values, np.ndarray):
+                X = X.values
+        except AttributeError:
+            pass
+
     # Iterate over columns for ndarrays.
     if isinstance(X, np.ndarray):
         X = X.T
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -559,6 +559,23 @@ def __getitem__(self, item):
     assert isinstance(xnew[0], np.ndarray)
 
 
+def test_reshape2d_pandas(pd):
+    # seperate to allow the rest of the tests to run if no pandas...
+    X = np.arange(30).reshape(10, 3)
+    x = pd.DataFrame(X, columns=["a", "b", "c"])
+    Xnew = cbook._reshape_2D(x, 'x')
+    # Need to check each row because _reshape_2D returns a list of arrays:
+    for x, xnew in zip(X.T, Xnew):
+        np.testing.assert_array_equal(x, xnew)
+
+    X = np.arange(30).reshape(10, 3)
+    x = pd.DataFrame(X, columns=["a", "b", "c"])
+    Xnew = cbook._reshape_2D(x, 'x')
+    # Need to check each row because _reshape_2D returns a list of arrays:
+    for x, xnew in zip(X.T, Xnew):
+        np.testing.assert_array_equal(x, xnew)
+
+
 def test_contiguous_regions():
     a, b, c = 3, 4, 5
     # Starts and ends with True
