diff --git a/yt/visualization/geo_plot_utils.py b/yt/visualization/geo_plot_utils.py index fc2c933e89e..db6716e42e8 100644 --- a/yt/visualization/geo_plot_utils.py +++ b/yt/visualization/geo_plot_utils.py @@ -1,4 +1,7 @@ -valid_transforms = {} +from types import FunctionType +from typing import Any, Dict, Optional, Tuple + +valid_transforms: Dict[str, FunctionType] = {} transform_list = [ "PlateCarree", @@ -40,7 +43,7 @@ def _func(*args, **kwargs): return _func -def get_mpl_transform(mpl_proj): +def get_mpl_transform(mpl_proj) -> Optional[FunctionType]: r"""This returns an instantiated transform function given a transform function name and arguments. @@ -67,11 +70,11 @@ def get_mpl_transform(mpl_proj): # check to see if mpl_proj is a string or tuple, and construct args and # kwargs to pass to cartopy function based on that. - key = None + key: Optional[str] = None + args: Tuple = () + kwargs: Dict[str, Any] = {} if isinstance(mpl_proj, str): key = mpl_proj - args = () - kwargs = {} instantiated_func = valid_transforms[key](*args, **kwargs) elif isinstance(mpl_proj, tuple): if len(mpl_proj) == 2: @@ -79,6 +82,12 @@ def get_mpl_transform(mpl_proj): kwargs = {} elif len(mpl_proj) == 3: key, args, kwargs = mpl_proj + else: + raise ValueError(f"Expected a tuple with len 2 or 3, received {mpl_proj}") + if not isinstance(key, str): + raise TypeError( + f"Expected a string a the first element in mpl_proj, got {key!r}" + ) instantiated_func = valid_transforms[key](*args, **kwargs) elif hasattr(mpl_proj, "globe"): # cartopy transforms have a globe method associated with them