Merge pull request #14289 from astrofrog/fix-reshape-2d-subclass

tacaswell · web-flow · commit 4edadf96576f · 2019-05-21T08:47:14.000-04:00
BUG: Fix performance regression when plotting values from Numpy array sub-classes
diff --git a/lib/matplotlib/cbook/__init__.py b/lib/matplotlib/cbook/__init__.py
@@ -1421,7 +1421,7 @@ def _reshape_2D(X, name):
     X = np.atleast_1d(X.T if isinstance(X, np.ndarray) else np.asarray(X))
     if len(X) == 0:
         return [[]]
-    if X.ndim == 1 and not isinstance(X[0], collections.abc.Iterable):
+    elif X.ndim == 1 and np.ndim(X[0]) == 0:
         # 1D array of scalars: directly return it.
         return [X]
     elif X.ndim in [1, 2]:
diff --git a/lib/matplotlib/tests/test_cbook.py b/lib/matplotlib/tests/test_cbook.py
@@ -492,8 +492,10 @@ def test_flatiter():
 
 
 def test_reshape2d():
+
     class dummy():
         pass
+
     xnew = cbook._reshape_2D([], 'x')
     assert np.shape(xnew) == (1, 0)
 
@@ -515,6 +517,41 @@ class dummy():
     xnew = cbook._reshape_2D(x, 'x')
     assert np.shape(xnew) == (5, 3)
 
+    # Now test with a list of lists with different lengths, which means the
+    # array will internally be converted to a 1D object array of lists
+    x = [[1, 2, 3], [3, 4], [2]]
+    xnew = cbook._reshape_2D(x, 'x')
+    assert isinstance(xnew, list)
+    assert isinstance(xnew[0], np.ndarray) and xnew[0].shape == (3,)
+    assert isinstance(xnew[1], np.ndarray) and xnew[1].shape == (2,)
+    assert isinstance(xnew[2], np.ndarray) and xnew[2].shape == (1,)
+
+    # We now need to make sure that this works correctly for Numpy subclasses
+    # where iterating over items can return subclasses too, which may be
+    # iterable even if they are scalars. To emulate this, we make a Numpy
+    # array subclass that returns Numpy 'scalars' when iterating or accessing
+    # values, and these are technically iterable if checking for example
+    # isinstance(x, collections.abc.Iterable).
+
+    class ArraySubclass(np.ndarray):
+
+        def __iter__(self):
+            for value in super().__iter__():
+                yield np.array(value)
+
+        def __getitem__(self, item):
+            return np.array(super().__getitem__(item))
+
+    v = np.arange(10, dtype=float)
+    x = ArraySubclass((10,), dtype=float, buffer=v.data)
+
+    xnew = cbook._reshape_2D(x, 'x')
+
+    # We check here that the array wasn't split up into many individual
+    # ArraySubclass, which is what used to happen due to a bug in _reshape_2D
+    assert len(xnew) == 1
+    assert isinstance(xnew[0], ArraySubclass)
+
 
 def test_contiguous_regions():
     a, b, c = 3, 4, 5
