From 76cb6d0261ca0b542131ccdb06ba6a3ddbc7581a Mon Sep 17 00:00:00 2001 From: Elliott Sales de Andrade Date: Mon, 14 Aug 2023 20:12:37 -0400 Subject: [PATCH] Backport PR #26504: TYP: Add overload to specify output of Colormap.__call__ when possible --- lib/matplotlib/colors.pyi | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/lib/matplotlib/colors.pyi b/lib/matplotlib/colors.pyi index 354250a7bd22..9bb1725f4f78 100644 --- a/lib/matplotlib/colors.pyi +++ b/lib/matplotlib/colors.pyi @@ -69,6 +69,15 @@ class Colormap: N: int colorbar_extend: bool def __init__(self, name: str, N: int = ...) -> None: ... + @overload + def __call__( + self, X: Sequence[float] | np.ndarray, alpha: ArrayLike | None = ..., bytes: bool = ... + ) -> np.ndarray: ... + @overload + def __call__( + self, X: float, alpha: float | None = ..., bytes: bool = ... + ) -> tuple[float, float, float, float]: ... + @overload def __call__( self, X: ArrayLike, alpha: ArrayLike | None = ..., bytes: bool = ... ) -> tuple[float, float, float, float] | np.ndarray: ...