Fix some new mypy errors

Author: mx-moth

src/emsarray/cli/commands/plot.py

@@ -33,7 +33,7 @@ def __init__(
     ):
         if default is None:
             default = {}
-        type = functools.partial(key_value, value_type=value_type)
+        type: Callable = functools.partial(key_value, value_type=value_type)
         super().__init__(
             option_strings, dest, default=default, type=type, **kwargs)
 

src/emsarray/plot.py

@@ -15,6 +15,7 @@
     import cartopy.crs
     from cartopy.feature import GSHHSFeature
     from cartopy.mpl import gridliner
+    from cartopy.mpl.geoaxes import GeoAxes
     from matplotlib import animation, patheffects
     from matplotlib.artist import Artist
     from matplotlib.axes import Axes
@@ -37,7 +38,7 @@
 CARTOPY_0_23 = CARTOPY_VERSION >= packaging.version.Version('0.23')
 
 
-def add_coast(axes: Axes, **kwargs: Any) -> None:
+def add_coast(axes: GeoAxes, **kwargs: Any) -> None:
     """
     Add coastlines to an :class:`~matplotlib.axes.Axes`.
     Some default styles are applied:
@@ -61,7 +62,7 @@ def add_coast(axes: Axes, **kwargs: Any) -> None:
     axes.add_feature(coast, **kwargs)
 
 
-def add_gridlines(axes: Axes, **kwargs: Any) -> gridliner.Gridliner:
+def add_gridlines(axes: GeoAxes, **kwargs: Any) -> gridliner.Gridliner:
     """
     Add some gridlines to the axes.
 
@@ -310,7 +311,7 @@ def plot_on_figure(
     if projection is None:
         projection = cartopy.crs.PlateCarree()
 
-    axes = figure.add_subplot(projection=projection)
+    axes: GeoAxes = figure.add_subplot(projection=projection)
     axes.set_aspect(aspect='equal', adjustable='datalim')
 
     if scalar is None and vector is None:

src/emsarray/transect.py

@@ -1,7 +1,7 @@
 import dataclasses
 from collections.abc import Hashable, Iterable
 from functools import cached_property
-from typing import Any, Callable, Generic, Optional, Union
+from typing import Any, Callable, Generic, Optional, Union, cast
 
 import cfunits
 import numpy
@@ -228,8 +228,8 @@ def _set_up_axis(self, variable: xarray.DataArray) -> tuple[str, Formatter]:
             # Use cfunits to normalize the units to their short symbol form.
             # EngFormatter will write 'k{unit}', 'G{unit}', etc
             # so unit symbols are required.
-            units = cfunits.Units(units).formatted()
-            formatter = EngFormatter(unit=units)
+            formatted_units = cfunits.Units(units).formatted()
+            formatter = EngFormatter(unit=formatted_units)
 
         return title, formatter
 
@@ -746,7 +746,7 @@ def _plot_on_figure(
             depth_limit_shallow = up(transect_dataset['depth_bounds'][depth_start])
         depth_limit_deep = down(transect_dataset['depth_bounds'][depth_stop])
 
-        axes = figure.subplots()
+        axes = cast(Axes, figure.subplots())
         x_title, x_formatter = self._set_up_axis(distance_bounds)
         y_title, y_formatter = self._set_up_axis(depth)
         axes.set_xlabel(x_title)
@@ -764,7 +764,8 @@ def _plot_on_figure(
         if title is not None:
             axes.set_title(title)
 
-        cmap = colormaps[cmap].copy()
+        if isinstance(cmap, str):
+            cmap = colormaps[cmap].copy()
         cmap.set_bad(ocean_floor_colour)
 
         if data_array.size != 0: