Use overloads for Normalize.__call__ and inverse, ensure NoNorm returns type consistent with superclass

ksunden · ksunden · commit 6385edc41606 · 2023-08-02T17:15:33.000-05:00
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -2095,9 +2095,13 @@ class NoNorm(Normalize):
     indices directly in a `~matplotlib.cm.ScalarMappable`.
     """
     def __call__(self, value, clip=None):
+        if np.iterable(value):
+            return np.ma.array(value)
         return value
 
     def inverse(self, value):
+        if np.iterable(value):
+            return np.ma.array(value)
         return value
 
 
diff --git a/lib/matplotlib/colors.pyi b/lib/matplotlib/colors.pyi
@@ -146,7 +146,17 @@ class Normalize:
     def clip(self, value: bool) -> None: ...
     @staticmethod
     def process_value(value: ArrayLike) -> tuple[np.ma.MaskedArray, bool]: ...
-    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ...
+    @overload
+    def __call__(self, value: float, clip: bool | None = ...) -> float: ...
+    @overload
+    def __call__(self, value: np.ndarray, clip: bool | None = ...) -> np.ma.MaskedArray: ...
+    @overload
+    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> ArrayLike: ...
+    @overload
+    def inverse(self, value: float) -> float: ...
+    @overload
+    def inverse(self, value: np.ndarray) -> np.ma.MaskedArray: ...
+    @overload
     def inverse(self, value: ArrayLike) -> ArrayLike: ...
     def autoscale(self, A: ArrayLike) -> None: ...
     def autoscale_None(self, A: ArrayLike) -> None: ...
@@ -161,8 +171,6 @@ class TwoSlopeNorm(Normalize):
     @vcenter.setter
     def vcenter(self, value: float) -> None: ...
     def autoscale_None(self, A: ArrayLike) -> None: ...
-    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ...
-    def inverse(self, value: ArrayLike) -> ArrayLike: ...
 
 class CenteredNorm(Normalize):
     def __init__(
@@ -240,8 +248,6 @@ class PowerNorm(Normalize):
         vmax: float | None = ...,
         clip: bool = ...,
     ) -> None: ...
-    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ...
-    def inverse(self, value: ArrayLike) -> ArrayLike: ...
 
 class BoundaryNorm(Normalize):
     boundaries: np.ndarray
@@ -256,12 +262,8 @@ class BoundaryNorm(Normalize):
         *,
         extend: Literal["neither", "both", "min", "max"] = ...
     ) -> None: ...
-    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ...
-    def inverse(self, value: ArrayLike) -> ArrayLike: ...
 
-class NoNorm(Normalize):
-    def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ...
-    def inverse(self, value: ArrayLike) -> ArrayLike: ...
+class NoNorm(Normalize): ...
 
 def rgb_to_hsv(arr: ArrayLike) -> np.ndarray: ...
 def hsv_to_rgb(hsv: ArrayLike) -> np.ndarray: ...
