Merge pull request #18870 from aitikgupta/scalarmappable-set_array

jklymak · web-flow · commit a18b4bfed7fb · 2021-05-13T07:16:30.000-07:00
Expand ScalarMappable.set_array to accept array-like inputs
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -361,12 +361,21 @@ 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
+            return
+
+        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/image.py b/lib/matplotlib/image.py
@@ -1346,8 +1346,7 @@ def make_image(self, renderer, magnification=1.0, unsampled=False):
 
     def set_data(self, A):
         """Set the image array."""
-        cm.ScalarMappable.set_array(self,
-                                    cbook.safe_masked_invalid(A, copy=True))
+        cm.ScalarMappable.set_array(self, A)
         self.stale = True
 
 
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -677,6 +677,22 @@ def test_collection_set_verts_array():
         assert np.array_equal(ap._codes, atp._codes)
 
 
+def test_collection_set_array():
+    vals = [*range(10)]
+
+    # Test set_array with list
+    c = Collection()
+    c.set_array(vals)
+
+    # Test set_array with wrong dtype
+    with pytest.raises(TypeError, match="^Image data of dtype"):
+        c.set_array("wrong_input")
+
+    # Test if array kwarg is copied
+    vals[5] = 45
+    assert np.not_equal(vals, c.get_array()).any()
+
+
 def test_blended_collection_autolim():
     a = [1, 2, 4]
     height = .2
