Override ScalarMappable set_array for Collection

aitikgupta · aitikgupta · commit 703343e67bfb · 2020-11-02T10:10:43.000+05:30
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -849,6 +849,23 @@ def set_alpha(self, alpha):
 
     set_alpha.__doc__ = artist.Artist._set_alpha_for_array.__doc__
 
+    def set_array(self, A):
+        """
+        Set the image array.
+
+        Parameters
+        ----------
+        A : array-like
+        """
+        # This method overrides the inherited
+        # cm.ScalarMappable.set_array method
+        self._A = cbook.safe_masked_invalid(A, copy=True)
+
+        if (self._A.dtype != np.uint8 and
+                not np.can_cast(self._A.dtype, float, "same_kind")):
+            raise TypeError("Image data of dtype {} cannot be converted to "
+                            "float".format(self._A.dtype))
+
     def get_linewidth(self):
         return self._linewidths
 
