removed shading=auto and changed tests

jeffreypaul15 · jeffreypaul15 · commit 2204dd198d92 · 2021-05-13T22:13:27.000+05:30
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -361,11 +361,11 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
 
     def set_array(self, A):
         """
-        Set the image array from numpy array *A*.
+        Set the image array from array-like *A*.
 
         Parameters
         ----------
-        A : ndarray or None
+        A : array-like or None
         """
         if A is None:
             self._A = None
@@ -375,14 +375,14 @@ def set_array(self, A):
             A_ = A.ravel()
             width, height = self._meshWidth, self._meshHeight
             if self._shading == 'flat':
-                if not ( len(A_) == width * height or 
-                len(A_) == (width + 1) * (height + 1)):
+                if not (len(A_) == width * height or
+                        len(A_) == (width + 1) * (height + 1)):
                     raise TypeError(
                         'Dimensions of A %s are incompatible with '
                         'X (%d) and/or Y (%d)' %
                         (A.shape, width, height))
             elif self._shading == 'nearest':
-                if not ( len(A_) == width * height:
+                if not (len(A_) == width * height):
                     raise TypeError(
                         'Dimensions of A %s are incompatible with '
                         'X (%d) and/or Y (%d)' %
@@ -394,6 +394,11 @@ def set_array(self, A):
                         'X (%d) and/or Y (%d)' %
                         (A.shape, width, height))
 
+        A = cbook.safe_masked_invalid(A, copy=True)
+        if not np.can_cast(A.dtype, float, "same_kind"):
+            raise TypeError(f"Image data of dtype {A.dtype} cannot be "
+                            "converted to float")
+
         self._A = A
 
     def get_array(self):
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -705,15 +705,16 @@ def test_quadmesh_set_array_validation():
     z = np.arange(9).reshape((3, 3))
     fig, ax = plt.subplots()
     coll = ax.pcolormesh(x, y, np.ones(z.shape))
-    z = np.arange(16).reshape((4, 4))
+    z = np.arange(25).reshape((5, 5))
 
-    with pytest.raises(TypeError, match=re.escape("Dimensions of A (4, 4) "
+    with pytest.raises(TypeError, match=re.escape("Dimensions of A (5, 5) "
                        "are incompatible with X (3) and/or Y (3)")):
         coll.set_array(z)
-    with pytest.raises(TypeError, match=re.escape("Dimensions of A (16,) "
+    with pytest.raises(TypeError, match=re.escape("Dimensions of A (25,) "
                        "are incompatible with X (3) and/or Y (3)")):
         coll.set_array(z.ravel())
 
+    z = np.arange(16).reshape((4, 4))
     fig, ax = plt.subplots()
     coll = ax.pcolormesh(x, y, np.ones(z.shape), shading='nearest')
     with pytest.raises(TypeError, match=re.escape("Dimensions of A (3, 3) "
