Skip to content

Commit

Permalink
TYP: fix type-check incompatibilities with matplotlib 3.8
Browse files Browse the repository at this point in the history
  • Loading branch information
neutrinoceros committed Sep 15, 2023
1 parent 8757cf2 commit f83e15c
Show file tree
Hide file tree
Showing 2 changed files with 25 additions and 20 deletions.
41 changes: 23 additions & 18 deletions yt/visualization/_handlers.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import sys
import weakref
from numbers import Real
from typing import Any, Literal, Optional, Union
from typing import TYPE_CHECKING, Any, Literal, Optional, Union

import matplotlib as mpl
import numpy as np
Expand All @@ -18,6 +18,24 @@
else:
from typing_extensions import TypeAlias

if TYPE_CHECKING:
# RGBColorType, RGBAColorType and ColorType are backported from matplotlib 3.8.0
RGBColorType = Union[tuple[float, float, float], str]
RGBAColorType = Union[
str, # "none" or "#RRGGBBAA"/"#RGBA" hex strings
tuple[float, float, float, float],
# 2 tuple (color, alpha) representations, not infinitely recursive
# RGBColorType includes the (str, float) tuple, even for RGBA strings
tuple[RGBColorType, float],
# (4-tuple, float) is odd, but accepted as the outer float overriding A of 4-tuple
tuple[tuple[float, float, float, float], float],
]

ColorType = Union[RGBColorType, RGBAColorType]

# this type alias is unique to the present module
ColormapInput: TypeAlias = Union[Colormap, str, None]


class NormHandler:
"""
Expand Down Expand Up @@ -398,16 +416,6 @@ def get_minmax(data):
return linthresh


BackgroundColor: TypeAlias = Union[
tuple[float, float, float, float],
# np.ndarray is only runtime-subscribtable since numpy 1.22
"np.ndarray[Any, Any]",
str,
None,
]
ColormapInput: TypeAlias = Union[Colormap, str, None]


class ColorbarHandler:
__slots__ = ("_draw_cbar", "_draw_minorticks", "_cmap", "_background_color")

Expand All @@ -423,7 +431,7 @@ def __init__(
self._draw_minorticks = draw_minorticks
self._cmap: Optional[Colormap] = None
self._set_cmap(cmap)
self._background_color: BackgroundColor = background_color
self._background_color: Optional["ColorType"] = background_color

@property
def draw_cbar(self) -> bool:
Expand Down Expand Up @@ -465,7 +473,7 @@ def _set_cmap(self, newval: ColormapInput) -> None:
self._cmap = newval
elif isinstance(newval, str):
self._cmap = mpl.colormaps[newval]
elif is_sequence(newval):
elif is_sequence(newval): # type: ignore[unreachable]
# tuple colormaps are from palettable (or brewer2mpl)
self._cmap = get_brewer_cmap(newval)
else:
Expand All @@ -475,14 +483,11 @@ def _set_cmap(self, newval: ColormapInput) -> None:
)

@property
def background_color(self) -> BackgroundColor:
def background_color(self) -> "ColorType":
return self._background_color or "white"

@background_color.setter
def background_color(self, newval: BackgroundColor) -> None:
# not attempting to constrain types here because
# down the line it really depends on matplotlib.axes.Axes.set_faceolor
# which is very type-flexibile
def background_color(self, newval: Optional["ColorType"]) -> None:
if newval is None:
self._background_color = self.cmap(0)
else:
Expand Down
4 changes: 2 additions & 2 deletions yt/visualization/base_plot_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import warnings
from abc import ABC
from io import BytesIO
from types import FunctionType
from typing import TYPE_CHECKING, Optional, TypedDict, Union

import matplotlib
Expand Down Expand Up @@ -37,6 +36,7 @@
from matplotlib.axes import Axes
from matplotlib.axis import Axis
from matplotlib.figure import Figure
from matplotlib.transforms import Transform

class FormatKwargs(TypedDict):
style: Literal["scientific"]
Expand Down Expand Up @@ -249,7 +249,7 @@ def __init__(
):
"""Initialize ImagePlotMPL class object"""

self._transform: Optional[FunctionType]
self._transform: Optional["Transform"]
setdefaultattr(self, "_transform", None)

self.colorbar_handler = colorbar_handler
Expand Down

0 comments on commit f83e15c

Please sign in to comment.