Allow passing axis kwargs to plot

doc/plotting.rst

@@ -743,12 +743,13 @@ This script will plot the air temperature on a map.
 
     air = xr.tutorial.open_dataset("air_temperature").air
 
-    ax = plt.axes(projection=ccrs.Orthographic(-80, 35))
-    air.isel(time=0).plot.contourf(ax=ax, transform=ccrs.PlateCarree())
-    ax.set_global()
+    p = air.isel(time=0).plot(
+        subplot_kws=dict(projection=ccrs.Orthographic(-80, 35), facecolor="gray"),
+        transform=ccrs.PlateCarree())
+    p.axes.set_global()
 
     @savefig plotting_maps_cartopy.png width=100%
-    ax.coastlines()
+    p.axes.coastlines()
 
 When faceting on maps, the projection can be transferred to the ``plot``
 function using the ``subplot_kws`` keyword. The axes for the subplots created

xarray/plot/plot.py

@@ -155,8 +155,7 @@ def plot(
         Relative tolerance used to determine if the indexes
         are uniformly spaced. Usually a small positive number.
     subplot_kws : dict, optional
-        Dictionary of keyword arguments for matplotlib subplots. Only applies
-        to FacetGrid plotting.
+        Dictionary of keyword arguments for matplotlib subplots.
     **kwargs : optional
         Additional keyword arguments to matplotlib
 
@@ -177,10 +176,10 @@ def plot(
 
     if ndims in [1, 2]:
         if row or col:
+            kwargs["subplot_kws"] = subplot_kws
             kwargs["row"] = row
             kwargs["col"] = col
             kwargs["col_wrap"] = col_wrap
-            kwargs["subplot_kws"] = subplot_kws
         if ndims == 1:
             plotfunc = line
             kwargs["hue"] = hue
@@ -190,6 +189,7 @@ def plot(
                 kwargs["hue"] = hue
             else:
                 plotfunc = pcolormesh
+                kwargs["subplot_kws"] = subplot_kws
     else:
         if row or col or hue:
             raise ValueError(error_msg)
@@ -548,8 +548,7 @@ def _plot2d(plotfunc):
         always infer intervals, unless the mesh is irregular and plotted on
         a map projection.
     subplot_kws : dict, optional
-        Dictionary of keyword arguments for matplotlib subplots. Only applies
-        to FacetGrid plotting.
+        Dictionary of keyword arguments for matplotlib subplots.
     cbar_ax : matplotlib Axes, optional
         Axes in which to draw the colorbar.
     cbar_kwargs : dict, optional
@@ -719,7 +718,10 @@ def newplotfunc(
                 "plt.imshow's `aspect` kwarg is not available " "in xarray"
             )
 
-        ax = get_axis(figsize, size, aspect, ax)
+        if subplot_kws is None:
+            subplot_kws = dict()
+        ax = get_axis(figsize, size, aspect, ax, **subplot_kws)
+
         primitive = plotfunc(
             xplt,
             yplt,

xarray/plot/utils.py

@@ -406,9 +406,12 @@ def _assert_valid_xy(darray, xy, name):
         raise ValueError(f"{name} must be one of None, '{valid_xy_str}'")
 
 
-def get_axis(figsize, size, aspect, ax):
-    import matplotlib as mpl
-    import matplotlib.pyplot as plt
+def get_axis(figsize=None, size=None, aspect=None, ax=None, **kwargs):
+    try:
+        import matplotlib as mpl
+        import matplotlib.pyplot as plt
+    except ImportError:
+        raise ImportError("matplotlib is required for plot.utils.get_axis")
 
     if figsize is not None:
         if ax is not None:
@@ -427,8 +430,11 @@ def get_axis(figsize, size, aspect, ax):
     elif aspect is not None:
         raise ValueError("cannot provide `aspect` argument without `size`")
 
+    if kwargs and ax is not None:
+        raise ValueError("cannot use subplot_kws with existing ax")
+
     if ax is None:
-        ax = plt.gca()
+        ax = plt.gca(**kwargs)
 
     return ax
 

xarray/tests/test_plot.py

@@ -15,6 +15,7 @@
     _build_discrete_cmap,
     _color_palette,
     _determine_cmap_params,
+    get_axis,
     label_from_attrs,
 )
 
@@ -2381,3 +2382,50 @@ def test_facetgrid_single_contour():
     ds["time"] = [0, 1]
 
     ds.plot.contour(col="time", levels=[4], colors=["k"])
+
+
+@requires_matplotlib
+@pytest.mark.parametrize("figsize", [[4, 4], None])
+@pytest.mark.parametrize("aspect", [4 / 3, None])
+@pytest.mark.parametrize("size", [200, None])
+def test_get_axis(figsize, aspect, size):
+    # test get_axis works with different args combinations
+    # and return the right type
+
+    ax = "something"
+    if figsize is not None:
+        # cannot provide both ax and figsize
+        with pytest.raises(ValueError, match="both `figsize` and `ax`"):
+            ax = get_axis(figsize, size, aspect, ax)
+    if figsize is None and size is not None:
+        # cannot provide both ax and size
+        with pytest.raises(ValueError, match="both `size` and `ax`"):
+            ax = get_axis(figsize, size, aspect, ax)
+
+    ax = None
+    if figsize is not None and size is not None:
+        # cannot provide both size and figsize
+        with pytest.raises(ValueError, match="both `figsize` and " "`size`"):
+            ax = get_axis(figsize, size, aspect, ax)
+
+    if figsize is None and aspect is not None and size is None:
+        # cannot provide aspect and size
+        with pytest.raises(ValueError, match="`aspect` argument without `size`"):
+            ax = get_axis(figsize, size, aspect, ax)
+
+    ax = plt.axes()
+    axtype = type(ax)
+    ax = get_axis(ax=ax)
+    assert isinstance(ax, axtype)
+    ax = get_axis()
+    assert isinstance(ax, axtype)
+
+    try:
+        import cartopy as ctpy  # type: ignore
+
+        kwargs = {"projection": ctpy.crs.PlateCarree()}
+        ax = get_axis(**kwargs)
+        assert isinstance(ax, ctpy.mpl.geoaxes.GeoAxesSubplot)
+    except ImportError:
+        # not testing cartopy
+        pass