improve typing of DataArray and Dataset reductions (#6746)

* improve typing of DataArray and Dataset reductions

* fix some mypy errors

* import Self from xarray.core.types

* fix typing errors

* fix remaining mypy error

* fix some typing issues

* add whats-new

* re-add type ignore

---------

Co-authored-by: Michael Niklas <[email protected]>

Author: rhkleijn

doc/whats-new.rst

@@ -26,7 +26,10 @@ New Features
   different collections of coordinates prior to assign them to a Dataset or
   DataArray (:pull:`8102`) at once.
   By `Benoît Bovy <https://github.com/benbovy>`_.
-- Provide `preferred_chunks` for data read from netcdf files (:issue:`1440`, :pull:`7948`)
+- Provide `preferred_chunks` for data read from netcdf files (:issue:`1440`, :pull:`7948`).
+  By `Martin Raspaud <https://github.com/mraspaud>`_.
+- Improved static typing of reduction methods (:pull:`6746`).
+  By `Richard Kleijn <https://github.com/rhkleijn>`_.
 
 Breaking changes
 ~~~~~~~~~~~~~~~~

xarray/core/_aggregations.py

@@ -8,7 +8,7 @@
 
 from xarray.core import duck_array_ops
 from xarray.core.options import OPTIONS
-from xarray.core.types import Dims
+from xarray.core.types import Dims, Self
 from xarray.core.utils import contains_only_chunked_or_numpy, module_available
 
 if TYPE_CHECKING:
@@ -30,7 +30,7 @@ def reduce(
         keep_attrs: bool | None = None,
         keepdims: bool = False,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         raise NotImplementedError()
 
     def count(
@@ -39,7 +39,7 @@ def count(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``count`` along some dimension(s).
 
@@ -111,7 +111,7 @@ def all(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``all`` along some dimension(s).
 
@@ -183,7 +183,7 @@ def any(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``any`` along some dimension(s).
 
@@ -256,7 +256,7 @@ def max(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``max`` along some dimension(s).
 
@@ -343,7 +343,7 @@ def min(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``min`` along some dimension(s).
 
@@ -430,7 +430,7 @@ def mean(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``mean`` along some dimension(s).
 
@@ -522,7 +522,7 @@ def prod(
         min_count: int | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``prod`` along some dimension(s).
 
@@ -629,7 +629,7 @@ def sum(
         min_count: int | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``sum`` along some dimension(s).
 
@@ -736,7 +736,7 @@ def std(
         ddof: int = 0,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``std`` along some dimension(s).
 
@@ -840,7 +840,7 @@ def var(
         ddof: int = 0,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``var`` along some dimension(s).
 
@@ -943,7 +943,7 @@ def median(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``median`` along some dimension(s).
 
@@ -1034,7 +1034,7 @@ def cumsum(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``cumsum`` along some dimension(s).
 
@@ -1127,7 +1127,7 @@ def cumprod(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> Dataset:
+    ) -> Self:
         """
         Reduce this Dataset's data by applying ``cumprod`` along some dimension(s).
 
@@ -1226,7 +1226,7 @@ def reduce(
         keep_attrs: bool | None = None,
         keepdims: bool = False,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         raise NotImplementedError()
 
     def count(
@@ -1235,7 +1235,7 @@ def count(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``count`` along some dimension(s).
 
@@ -1301,7 +1301,7 @@ def all(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``all`` along some dimension(s).
 
@@ -1367,7 +1367,7 @@ def any(
         *,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``any`` along some dimension(s).
 
@@ -1434,7 +1434,7 @@ def max(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``max`` along some dimension(s).
 
@@ -1513,7 +1513,7 @@ def min(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``min`` along some dimension(s).
 
@@ -1592,7 +1592,7 @@ def mean(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``mean`` along some dimension(s).
 
@@ -1676,7 +1676,7 @@ def prod(
         min_count: int | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``prod`` along some dimension(s).
 
@@ -1773,7 +1773,7 @@ def sum(
         min_count: int | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``sum`` along some dimension(s).
 
@@ -1870,7 +1870,7 @@ def std(
         ddof: int = 0,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``std`` along some dimension(s).
 
@@ -1964,7 +1964,7 @@ def var(
         ddof: int = 0,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``var`` along some dimension(s).
 
@@ -2057,7 +2057,7 @@ def median(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``median`` along some dimension(s).
 
@@ -2140,7 +2140,7 @@ def cumsum(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``cumsum`` along some dimension(s).
 
@@ -2229,7 +2229,7 @@ def cumprod(
         skipna: bool | None = None,
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> DataArray:
+    ) -> Self:
         """
         Reduce this DataArray's data by applying ``cumprod`` along some dimension(s).
 

xarray/core/computation.py

@@ -1402,14 +1402,14 @@ def _cov_corr(
     ) / (valid_count)
 
     if method == "cov":
-        return cov  # type: ignore[return-value]
+        return cov
 
     else:
         # compute std + corr
         da_a_std = da_a.std(dim=dim)
         da_b_std = da_b.std(dim=dim)
         corr = cov / (da_a_std * da_b_std)
-        return corr  # type: ignore[return-value]
+        return corr
 
 
 def cross(

xarray/core/rolling.py

@@ -475,9 +475,8 @@ def reduce(
             obj, rolling_dim, keep_attrs=keep_attrs, fill_value=fillna
         )
 
-        result = windows.reduce(
-            func, dim=list(rolling_dim.values()), keep_attrs=keep_attrs, **kwargs
-        )
+        dim = list(rolling_dim.values())
+        result = windows.reduce(func, dim=dim, keep_attrs=keep_attrs, **kwargs)
 
         # Find valid windows based on count.
         counts = self._counts(keep_attrs=False)
@@ -494,14 +493,15 @@ def _counts(self, keep_attrs: bool | None) -> DataArray:
         # array is faster to be reduced than object array.
         # The use of skipna==False is also faster since it does not need to
         # copy the strided array.
+        dim = list(rolling_dim.values())
         counts = (
             self.obj.notnull(keep_attrs=keep_attrs)
             .rolling(
                 {d: w for d, w in zip(self.dim, self.window)},
                 center={d: self.center[i] for i, d in enumerate(self.dim)},
             )
             .construct(rolling_dim, fill_value=False, keep_attrs=keep_attrs)
-            .sum(dim=list(rolling_dim.values()), skipna=False, keep_attrs=keep_attrs)
+            .sum(dim=dim, skipna=False, keep_attrs=keep_attrs)
         )
         return counts
 

xarray/tests/test_groupby.py

@@ -728,7 +728,7 @@ def test_groupby_dataset_iter() -> None:
 def test_groupby_dataset_errors() -> None:
     data = create_test_data()
     with pytest.raises(TypeError, match=r"`group` must be"):
-        data.groupby(np.arange(10))  # type: ignore
+        data.groupby(np.arange(10))  # type: ignore[arg-type,unused-ignore]
     with pytest.raises(ValueError, match=r"length does not match"):
         data.groupby(data["dim1"][:3])
     with pytest.raises(TypeError, match=r"`group` must be"):

xarray/util/generate_aggregations.py

@@ -27,7 +27,7 @@
 
 from xarray.core import duck_array_ops
 from xarray.core.options import OPTIONS
-from xarray.core.types import Dims
+from xarray.core.types import Dims, Self
 from xarray.core.utils import contains_only_chunked_or_numpy, module_available
 
 if TYPE_CHECKING:
@@ -50,7 +50,7 @@ def reduce(
         keep_attrs: bool | None = None,
         keepdims: bool = False,
         **kwargs: Any,
-    ) -> {obj}:
+    ) -> Self:
         raise NotImplementedError()"""
 
 GROUPBY_PREAMBLE = """
@@ -108,7 +108,7 @@ def {method}(
         *,{extra_kwargs}
         keep_attrs: bool | None = None,
         **kwargs: Any,
-    ) -> {obj}:
+    ) -> Self:
         """
         Reduce this {obj}'s data by applying ``{method}`` along some dimension(s).