From 47d155d7100bb0eb6e0c5eacd75f8e1ac1e0c3b5 Mon Sep 17 00:00:00 2001 From: Kyle Sunden Date: Tue, 1 Aug 2023 19:59:51 -0500 Subject: [PATCH 1/2] Adjust type hint of Norm.__call__ to return masked array xref pydata/xarray#8030 We were perhaps overly permissive in the return type, I'm pretty sure we always return Masked Array, so may as well say so, so that downstream can rely on that (which they already are, to be clear...) --- lib/matplotlib/colors.pyi | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/matplotlib/colors.pyi b/lib/matplotlib/colors.pyi index e222077cda14..49ffdf66753d 100644 --- a/lib/matplotlib/colors.pyi +++ b/lib/matplotlib/colors.pyi @@ -146,7 +146,7 @@ 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 = ...) -> ArrayLike: ... + def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ... def inverse(self, value: ArrayLike) -> ArrayLike: ... def autoscale(self, A: ArrayLike) -> None: ... def autoscale_None(self, A: ArrayLike) -> None: ... @@ -161,7 +161,7 @@ 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 = ...) -> ArrayLike: ... + def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ... def inverse(self, value: ArrayLike) -> ArrayLike: ... class CenteredNorm(Normalize): @@ -240,7 +240,7 @@ class PowerNorm(Normalize): vmax: float | None = ..., clip: bool = ..., ) -> None: ... - def __call__(self, value: ArrayLike, clip: bool | None = ...) -> ArrayLike: ... + def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ... def inverse(self, value: ArrayLike) -> ArrayLike: ... class BoundaryNorm(Normalize): @@ -256,11 +256,11 @@ class BoundaryNorm(Normalize): *, extend: Literal["neither", "both", "min", "max"] = ... ) -> None: ... - def __call__(self, value: ArrayLike, clip: bool | None = ...) -> ArrayLike: ... + 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 = ...) -> ArrayLike: ... + def __call__(self, value: ArrayLike, clip: bool | None = ...) -> np.ma.MaskedArray: ... def inverse(self, value: ArrayLike) -> ArrayLike: ... def rgb_to_hsv(arr: ArrayLike) -> np.ndarray: ... From 6385edc416066902ec662b8af35363453b43f40f Mon Sep 17 00:00:00 2001 From: Kyle Sunden Date: Wed, 2 Aug 2023 17:01:39 -0500 Subject: [PATCH 2/2] Use overloads for Normalize.__call__ and inverse, ensure NoNorm returns type consistent with superclass --- lib/matplotlib/colors.py | 4 ++++ lib/matplotlib/colors.pyi | 22 ++++++++++++---------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py index d925230584ca..b348560d1775 100644 --- 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 index 49ffdf66753d..758890bda2f8 100644 --- 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: ...