better use of union for colormap types

trygvrad · trygvrad · commit 37f1e7eef2ae · 2024-08-01T21:24:11.000+02:00
Defines ColorMapType as Union[colors.Colormap, colors.BivarColormap, colors.MultivarColormap] in typing.py
diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py
@@ -27,11 +27,7 @@
 from matplotlib._cm_listed import cmaps as cmaps_listed
 from matplotlib._cm_multivar import cmap_families as multivar_cmaps
 from matplotlib._cm_bivar import cmaps as bivar_cmaps
-from typing import Union
-
-ColorMapType = Union[colors.Colormap,
-                     colors.BivarColormap,
-                     colors.MultivarColormap]
+from .typing import ColorMapType
 
 _LUTSIZE = mpl.rcParams['image.lut']
 
@@ -227,9 +223,7 @@ def get_cmap(self, cmap):
             return self[mpl.rcParams["image.cmap"]]
 
         # if the user passed in a valid colormap type, simply return it
-        if isinstance(cmap, (colors.Colormap,
-                             colors.BivarColormap,
-                             colors.MultivarColormap)):
+        if isinstance(cmap, ColorMapType):
             return cmap
         if isinstance(cmap, str):
             _api.check_in_list(sorted(_colormaps), cmap=cmap)
diff --git a/lib/matplotlib/cm.pyi b/lib/matplotlib/cm.pyi
@@ -5,10 +5,7 @@ from matplotlib.colorbar import Colorbar
 import numpy as np
 from numpy.typing import ArrayLike
 from typing import Union
-
-ColorMapType = Union[colors.Colormap,
-                     colors.BivarColormap,
-                     colors.MultivarColormap]
+from .typing import ColorMapType
 
 class ColormapRegistry(Mapping[str, ColorMapType]):
     def __init__(self, cmaps: Mapping[str, ColorMapType]) -> None: ...
diff --git a/lib/matplotlib/typing.py b/lib/matplotlib/typing.py
@@ -13,7 +13,7 @@
 import pathlib
 from typing import Any, Literal, TypeAlias, TypeVar
 
-from . import path
+from . import path, colors
 from ._enums import JoinStyle, CapStyle
 from .markers import MarkerStyle
 
@@ -34,6 +34,7 @@
 RGBAColourType: TypeAlias = RGBAColorType
 ColourType: TypeAlias = ColorType
 
+ColorMapType = Union[colors.Colormap, colors.BivarColormap, colors.MultivarColormap]
 LineStyleType: TypeAlias = str | tuple[float, Sequence[float]]
 DrawStyleType: TypeAlias = Literal["default", "steps", "steps-pre", "steps-mid",
                                    "steps-post"]
