simplify to_rgba() by extracting the part relating to RGBA data

trygvrad · trygvrad · commit d13000627c9e · 2024-08-09T12:16:41.000+02:00
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -1446,7 +1446,7 @@ def set_array(self, A):
         A : array-like or None
             The values that are mapped to colors.
 
-            The base class `.VectorMappable` does not make any assumptions on
+            The base class `.ColorizingArtist` does not make any assumptions on
             the dimensionality and shape of the value array *A*.
         """
         if A is None:
@@ -1466,7 +1466,7 @@ def get_array(self):
         """
         Return the array of values, that are mapped to colors.
 
-        The base class `.VectorMappable` does not make any assumptions on
+        The base class `.ColorizingArtist` does not make any assumptions on
         the dimensionality and shape of the array.
         """
         return self._A
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -279,7 +279,7 @@ class ScalarMappable(colorizer.ColorizerShim):
     """
     A mixin class to map one or multiple sets of scalar data to RGBA.
 
-    The VectorMappable applies data normalization before returning RGBA colors
+    The ScalarMappable applies data normalization before returning RGBA colors
     from the given `~matplotlib.colors.Colormap`, `~matplotlib.colors.BivarColormap`,
     or `~matplotlib.colors.MultivarColormap`.
     """
diff --git a/lib/matplotlib/colorizer.py b/lib/matplotlib/colorizer.py
@@ -24,8 +24,8 @@
 class Colorizer():
     """
     Class that holds the data to color pipeline
-    accessible via `.to_rgba(A)` and executed via
-    the `.norm` and `.cmap` attributes.
+    accessible via `Colorizer.to_rgba(A)` and executed via
+    the `Colorizer.norm` and `Colorizer.cmap` attributes.
     """
     def __init__(self, cmap=None, norm=None):
 
@@ -125,56 +125,59 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True):
 
         """
         # First check for special case, image input:
-        # First check for special case, image input:
-        try:
-            if x.ndim == 3:
-                if x.shape[2] == 3:
-                    if alpha is None:
-                        alpha = 1
-                    if x.dtype == np.uint8:
-                        alpha = np.uint8(alpha * 255)
-                    m, n = x.shape[:2]
-                    xx = np.empty(shape=(m, n, 4), dtype=x.dtype)
-                    xx[:, :, :3] = x
-                    xx[:, :, 3] = alpha
-                elif x.shape[2] == 4:
-                    xx = x
-                else:
-                    raise ValueError("Third dimension must be 3 or 4")
-                if xx.dtype.kind == 'f':
-                    # If any of R, G, B, or A is nan, set to 0
-                    if np.any(nans := np.isnan(x)):
-                        if x.shape[2] == 4:
-                            xx = xx.copy()
-                        xx[np.any(nans, axis=2), :] = 0
-
-                    if norm and (xx.max() > 1 or xx.min() < 0):
-                        raise ValueError("Floating point image RGB values "
-                                         "must be in the 0..1 range.")
-                    if bytes:
-                        xx = (xx * 255).astype(np.uint8)
-                elif xx.dtype == np.uint8:
-                    if not bytes:
-                        xx = xx.astype(np.float32) / 255
-                else:
-                    raise ValueError("Image RGB array must be uint8 or "
-                                     "floating point; found %s" % xx.dtype)
-                # Account for any masked entries in the original array
-                # If any of R, G, B, or A are masked for an entry, we set alpha to 0
-                if np.ma.is_masked(x):
-                    xx[np.any(np.ma.getmaskarray(x), axis=2), 3] = 0
-                return xx
-        except AttributeError:
-            # e.g., x is not an ndarray; so try mapping it
-            pass
-
-        # This is the normal case, mapping a scalar array:
+        if isinstance(x, np.ndarray) and x.ndim == 3:
+            return self._pass_image_data(x, alpha, bytes, norm)
+
+        # Otherwise run norm -> colormap pipeline
         x = ma.asarray(x)
         if norm:
             x = self.norm(x)
         rgba = self.cmap(x, alpha=alpha, bytes=bytes)
         return rgba
 
+    @staticmethod
+    def _pass_image_data(x, alpha=None, bytes=False, norm=True):
+        """
+        Helper function to pass ndarray of shape (...,3) or (..., 4)
+        through `to_rgba()`, see `to_rgba()` for docstring.
+        """
+        if x.shape[2] == 3:
+            if alpha is None:
+                alpha = 1
+            if x.dtype == np.uint8:
+                alpha = np.uint8(alpha * 255)
+            m, n = x.shape[:2]
+            xx = np.empty(shape=(m, n, 4), dtype=x.dtype)
+            xx[:, :, :3] = x
+            xx[:, :, 3] = alpha
+        elif x.shape[2] == 4:
+            xx = x
+        else:
+            raise ValueError("Third dimension must be 3 or 4")
+        if xx.dtype.kind == 'f':
+            # If any of R, G, B, or A is nan, set to 0
+            if np.any(nans := np.isnan(x)):
+                if x.shape[2] == 4:
+                    xx = xx.copy()
+                xx[np.any(nans, axis=2), :] = 0
+
+            if norm and (xx.max() > 1 or xx.min() < 0):
+                raise ValueError("Floating point image RGB values "
+                                 "must be in the 0..1 range.")
+            if bytes:
+                xx = (xx * 255).astype(np.uint8)
+        elif xx.dtype == np.uint8:
+            if not bytes:
+                xx = xx.astype(np.float32) / 255
+        else:
+            raise ValueError("Image RGB array must be uint8 or "
+                             "floating point; found %s" % xx.dtype)
+        # Account for any masked entries in the original array
+        # If any of R, G, B, or A are masked for an entry, we set alpha to 0
+        if np.ma.is_masked(x):
+            xx[np.any(np.ma.getmaskarray(x), axis=2), 3] = 0
+        return xx
+
     def normalize(self, x):
         """
         Normalize the data in x.
