Fix passing ax in quick plot methods

monet/monet_accessor.py

@@ -975,6 +975,7 @@ def quick_imshow(self, map_kws=None, roll_dateline=False, **kwargs):
         import cartopy.crs as ccrs
         import matplotlib.pyplot as plt
         import seaborn as sns
+        from cartopy.mpl.geoaxes import GeoAxes
 
         from .plots import _dynamic_fig_size
         from .plots.mapgen import draw_map
@@ -1001,6 +1002,10 @@ def quick_imshow(self, map_kws=None, roll_dateline=False, **kwargs):
             kwargs.pop("transform", None)
         if "ax" not in kwargs:
             ax = draw_map(**map_kws)
+        else:
+            ax = kwargs.pop("ax", None)
+            if not isinstance(ax, GeoAxes):
+                raise TypeError("`ax` should be a Cartopy GeoAxes instance")
         _set_outline_patch_alpha(ax)
         if roll_dateline:
             _ = (
@@ -1035,6 +1040,7 @@ def quick_map(self, map_kws=None, roll_dateline=False, **kwargs):
         import cartopy.crs as ccrs
         import matplotlib.pyplot as plt
         import seaborn as sns
+        from cartopy.mpl.geoaxes import GeoAxes
 
         from .plots import _dynamic_fig_size
         from .plots.mapgen import draw_map
@@ -1056,6 +1062,10 @@ def quick_map(self, map_kws=None, roll_dateline=False, **kwargs):
         transform = kwargs.pop("transform", crs_p)
         if "ax" not in kwargs:
             ax = draw_map(**map_kws)
+        else:
+            ax = kwargs.pop("ax", None)
+            if not isinstance(ax, GeoAxes):
+                raise TypeError("`ax` should be a Cartopy GeoAxes instance")
         _set_outline_patch_alpha(ax)
         if roll_dateline:
             _ = da.roll(x=int(len(da.x) / 2), roll_coords=True).plot(
@@ -1088,6 +1098,7 @@ def quick_contourf(self, map_kws=None, roll_dateline=False, **kwargs):
         import cartopy.crs as ccrs
         import matplotlib.pyplot as plt
         import seaborn as sns
+        from cartopy.mpl.geoaxes import GeoAxes
 
         from monet.plots import _dynamic_fig_size
         from monet.plots.mapgen import draw_map
@@ -1113,6 +1124,10 @@ def quick_contourf(self, map_kws=None, roll_dateline=False, **kwargs):
             kwargs.pop("transform", None)
         if "ax" not in kwargs:
             ax = draw_map(**map_kws)
+        else:
+            ax = kwargs.pop("ax", None)
+            if not isinstance(ax, GeoAxes):
+                raise TypeError("`ax` should be a Cartopy GeoAxes instance")
         _set_outline_patch_alpha(ax)
         if roll_dateline:
             _ = da.roll(x=int(len(da.x) / 2), roll_coords=True).plot.contourf(

monet/plots/__init__.py

@@ -202,6 +202,7 @@ def sp_scatter_bias(
 
 
 def _set_outline_patch_alpha(ax, alpha=0):
+    """For :class:`cartopy.mpl.geoaxes.GeoAxes`"""
     for f in [
         lambda alpha: ax.axes.outline_patch.set_alpha(alpha),
         lambda alpha: ax.outline_patch.set_alpha(alpha),

tests/test_plot.py

@@ -1,4 +1,6 @@
 import cartopy
+import cartopy.crs as ccrs
+import matplotlib.pyplot as plt
 import pytest
 import xarray as xr
 from packaging.version import Version
@@ -11,13 +13,28 @@
 
 
 @pytest.mark.parametrize("which", ["imshow", "map", "contourf"])
-def test_quick_(which):
+def test_quick(which):
     getattr(da.monet, f"quick_{which}")()
 
 
-if __name__ == "__main__":
-    import matplotlib.pyplot as plt
+@pytest.mark.parametrize("which", ["imshow", "map", "contourf"])
+def test_quick_with_ax(which):
+    _, ax = plt.subplots()
+
+    with pytest.raises(TypeError, match="`ax` should be a Cartopy GeoAxes instance"):
+        getattr(da.monet, f"quick_{which}")(ax=ax)
+
+
+@pytest.mark.parametrize("which", ["imshow", "map", "contourf"])
+def test_quick_with_cartopy_ax(which):
+    proj = tran = ccrs.PlateCarree()
+
+    _, ax = plt.subplots(subplot_kw=dict(projection=proj))
+
+    getattr(da.monet, f"quick_{which}")(ax=ax, transform=tran)
 
-    test_quick_("map")
+
+if __name__ == "__main__":
+    test_quick("map")
 
     plt.show()