Merge pull request #15393 from anntzer/reshape

QuLogic · web-flow · commit aaedb1ab7a39 · 2020-04-01T16:58:15.000-04:00
Cleanup shape manipulations.
diff --git a/examples/pyplots/boxplot_demo_pyplot.py b/examples/pyplots/boxplot_demo_pyplot.py
@@ -65,16 +65,14 @@
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
 d2 = np.concatenate((spread, center, flier_high, flier_low))
-data.shape = (-1, 1)
-d2.shape = (-1, 1)
 
 ###############################################################################
 # Making a 2-D array only works if all the columns are the
 # same length.  If they are not, then use a list instead.
 # This is actually more efficient because boxplot converts
 # a 2-D array into a list of vectors internally anyway.
 
-data = [data, d2, d2[::2, 0]]
+data = [data, d2, d2[::2]]
 fig7, ax7 = plt.subplots()
 ax7.set_title('Multiple Samples with Different sizes')
 ax7.boxplot(data)
diff --git a/examples/statistics/boxplot_demo.py b/examples/statistics/boxplot_demo.py
@@ -60,13 +60,11 @@
 flier_high = np.random.rand(10) * 100 + 100
 flier_low = np.random.rand(10) * -100
 d2 = np.concatenate((spread, center, flier_high, flier_low))
-data.shape = (-1, 1)
-d2.shape = (-1, 1)
 # Making a 2-D array only works if all the columns are the
 # same length.  If they are not, then use a list instead.
 # This is actually more efficient because boxplot converts
 # a 2-D array into a list of vectors internally anyway.
-data = [data, d2, d2[::2, 0]]
+data = [data, d2, d2[::2]]
 
 # Multiple box plots on one Axes
 fig, ax = plt.subplots()
diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py
@@ -1069,7 +1069,7 @@ def set_data(self, x, y, A):
             raise TypeError("3D arrays must have three (RGB) "
                             "or four (RGBA) color components")
         if A.ndim == 3 and A.shape[2] == 1:
-            A.shape = A.shape[0:2]
+            A = A.squeeze(axis=-1)
         self._A = A
         self._Ax = x
         self._Ay = y
@@ -1230,7 +1230,7 @@ def set_data(self, x, y, A):
         if A.ndim not in [2, 3]:
             raise ValueError("A must be 2D or 3D")
         if A.ndim == 3 and A.shape[2] == 1:
-            A.shape = A.shape[:2]
+            A = A.squeeze(axis=-1)
         self._is_grayscale = False
         if A.ndim == 3:
             if A.shape[2] in [3, 4]:
diff --git a/lib/matplotlib/tests/test_image.py b/lib/matplotlib/tests/test_image.py
@@ -657,17 +657,6 @@ def test_jpeg_alpha():
     assert corner_pixel == (254, 0, 0)
 
 
-def test_nonuniformimage_setdata():
-    ax = plt.gca()
-    im = NonUniformImage(ax)
-    x = np.arange(3, dtype=float)
-    y = np.arange(4, dtype=float)
-    z = np.arange(12, dtype=float).reshape((4, 3))
-    im.set_data(x, y, z)
-    x[0] = y[0] = z[0, 0] = 9.9
-    assert im._A[0, 0] == im._Ax[0] == im._Ay[0] == 0, 'value changed'
-
-
 def test_axesimage_setdata():
     ax = plt.gca()
     im = AxesImage(ax)
@@ -686,15 +675,20 @@ def test_figureimage_setdata():
     assert im._A[0, 0] == 0, 'value changed'
 
 
-def test_pcolorimage_setdata():
+@pytest.mark.parametrize(
+    "image_cls,x,y,a", [
+        (NonUniformImage,
+         np.arange(3.), np.arange(4.), np.arange(12.).reshape((4, 3))),
+        (PcolorImage,
+         np.arange(3.), np.arange(4.), np.arange(6.).reshape((3, 2))),
+    ])
+def test_setdata_xya(image_cls, x, y, a):
     ax = plt.gca()
-    im = PcolorImage(ax)
-    x = np.arange(3, dtype=float)
-    y = np.arange(4, dtype=float)
-    z = np.arange(6, dtype=float).reshape((3, 2))
-    im.set_data(x, y, z)
-    x[0] = y[0] = z[0, 0] = 9.9
+    im = image_cls(ax)
+    im.set_data(x, y, a)
+    x[0] = y[0] = a[0, 0] = 9.9
     assert im._A[0, 0] == im._Ax[0] == im._Ay[0] == 0, 'value changed'
+    im.set_data(x, y, a.reshape((*a.shape, -1)))  # Just a smoketest.
 
 
 def test_minimized_rasterized():
