Add typing to plot methods

doc/whats-new.rst

@@ -70,6 +70,8 @@ Bug fixes
   By `Deepak Cherian <https://github.com/dcherian>`_.
 - ``Dataset.encoding['source']`` now exists when reading from a Path object (:issue:`5888`, :pull:`6974`)
   By `Thomas Coleman <https://github.com/ColemanTom>`_.
+- Add initial static typing to plot accessors (:issue:`6949`, :pull:`7052`).
+  By `Michael Niklas <https://github.com/headtr1ck>`_.
 
 Documentation
 ~~~~~~~~~~~~~

setup.cfg

@@ -203,6 +203,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-matplotlib.*]
 ignore_missing_imports = True
+[mypy-mpl_toolkits.*]
+ignore_missing_imports = True
 [mypy-Nio.*]
 ignore_missing_imports = True
 [mypy-nc_time_axis.*]

xarray/core/dataarray.py

@@ -22,7 +22,7 @@
 
 from ..coding.calendar_ops import convert_calendar, interp_calendar
 from ..coding.cftimeindex import CFTimeIndex
-from ..plot.plot import _PlotMethods
+from ..plot.plot import DataArrayPlotAccessor
 from ..plot.utils import _get_units_from_attrs
 from . import alignment, computation, dtypes, indexing, ops, utils
 from ._reductions import DataArrayReductions
@@ -3575,7 +3575,7 @@ def _inplace_binary_op(self: T_DataArray, other: Any, f: Callable) -> T_DataArra
     def _copy_attrs_from(self, other: DataArray | Dataset | Variable) -> None:
         self.attrs = other.attrs
 
-    plot = utils.UncachedAccessor(_PlotMethods)
+    plot = utils.UncachedAccessor(DataArrayPlotAccessor)
 
     def _title_for_slice(self, truncate: int = 50) -> str:
         """

xarray/core/dataset.py

@@ -35,7 +35,7 @@
 
 from ..coding.calendar_ops import convert_calendar, interp_calendar
 from ..coding.cftimeindex import CFTimeIndex, _parse_array_of_cftime_strings
-from ..plot.dataset_plot import _Dataset_PlotMethods
+from ..plot.dataset_plot import DatasetPlotAccessor
 from . import alignment
 from . import dtypes as xrdtypes
 from . import duck_array_ops, formatting, formatting_html, ops, utils
@@ -7198,7 +7198,7 @@ def real(self: T_Dataset) -> T_Dataset:
     def imag(self: T_Dataset) -> T_Dataset:
         return self.map(lambda x: x.imag, keep_attrs=True)
 
-    plot = utils.UncachedAccessor(_Dataset_PlotMethods)
+    plot = utils.UncachedAccessor(DatasetPlotAccessor)
 
     def filter_by_attrs(self: T_Dataset, **kwargs) -> T_Dataset:
         """Returns a ``Dataset`` with variables that match specific conditions.

xarray/core/types.py

@@ -114,6 +114,8 @@
 CoarsenBoundaryOptions = Literal["exact", "trim", "pad"]
 SideOptions = Literal["left", "right"]
 
+MPLScaleOptions = Literal["linear", "symlog", "log", "logit", None]
+
 # TODO: Wait until mypy supports recursive objects in combination with typevars
 _T = TypeVar("_T")
 NestedSequence = Union[

xarray/plot/dataset_plot.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import functools
+from typing import TYPE_CHECKING, NoReturn
 
 import numpy as np
 import pandas as pd
@@ -16,6 +17,9 @@
     get_axis,
 )
 
+if TYPE_CHECKING:
+    from ..core.dataset import Dataset
+
 
 def _infer_scatter_data(ds, x, y, hue, markersize, size_norm, size_mapping=None):
 
@@ -48,22 +52,6 @@ def _infer_scatter_data(ds, x, y, hue, markersize, size_norm, size_mapping=None)
     return data
 
 
-class _Dataset_PlotMethods:
-    """
-    Enables use of xarray.plot functions as attributes on a Dataset.
-    For example, Dataset.plot.scatter
-    """
-
-    def __init__(self, dataset):
-        self._ds = dataset
-
-    def __call__(self, *args, **kwargs):
-        raise ValueError(
-            "Dataset.plot cannot be called directly. Use "
-            "an explicit plot method, e.g. ds.plot.scatter(...)"
-        )
-
-
 def _dsplot(plotfunc):
     commondoc = """
     Parameters
@@ -299,55 +287,10 @@ def newplotfunc(
 
         return primitive
 
-    @functools.wraps(newplotfunc)
-    def plotmethod(
-        _PlotMethods_obj,
-        x=None,
-        y=None,
-        u=None,
-        v=None,
-        hue=None,
-        hue_style=None,
-        col=None,
-        row=None,
-        ax=None,
-        figsize=None,
-        col_wrap=None,
-        sharex=True,
-        sharey=True,
-        aspect=None,
-        size=None,
-        subplot_kws=None,
-        add_guide=None,
-        cbar_kwargs=None,
-        cbar_ax=None,
-        vmin=None,
-        vmax=None,
-        norm=None,
-        infer_intervals=None,
-        center=None,
-        levels=None,
-        robust=None,
-        colors=None,
-        extend=None,
-        cmap=None,
-        **kwargs,
-    ):
-        """
-        The method should have the same signature as the function.
-
-        This just makes the method work on Plotmethods objects,
-        and passes all the other arguments straight through.
-        """
-        allargs = locals()
-        allargs["ds"] = _PlotMethods_obj._ds
-        allargs.update(kwargs)
-        for arg in ["_PlotMethods_obj", "newplotfunc", "kwargs"]:
-            del allargs[arg]
-        return newplotfunc(**allargs)
-
-    # Add to class _PlotMethods
-    setattr(_Dataset_PlotMethods, plotmethod.__name__, plotmethod)
+    # we want to actually expose the signature of newplotfunc
+    # and not the copied **kwargs from the plotfunc which
+    # functools.wraps adds, so delete the wrapped attr
+    del newplotfunc.__wrapped__
 
     return newplotfunc
 
@@ -497,3 +440,33 @@ def streamplot(ds, x, y, ax, u, v, **kwargs):
 
     # Return .lines so colorbar creation works properly
     return hdl.lines
+
+
+class DatasetPlotAccessor:
+    """
+    Enables use of xarray.plot functions as attributes on a Dataset.
+    For example, Dataset.plot.scatter
+    """
+
+    _ds: Dataset
+
+    def __init__(self, dataset: Dataset) -> None:
+        self._ds = dataset
+
+    def __call__(self, *args, **kwargs) -> NoReturn:
+        raise ValueError(
+            "Dataset.plot cannot be called directly. Use "
+            "an explicit plot method, e.g. ds.plot.scatter(...)"
+        )
+
+    @functools.wraps(scatter)
+    def scatter(self, *args, **kwargs):
+        return scatter(self._ds, *args, **kwargs)
+
+    @functools.wraps(quiver)
+    def quiver(self, *args, **kwargs):
+        return quiver(self._ds, *args, **kwargs)
+
+    @functools.wraps(streamplot)
+    def streamplot(self, *args, **kwargs):
+        return streamplot(self._ds, *args, **kwargs)