Skip to content

Commit

Permalink
improve typing of DataArray and Dataset reductions
Browse files Browse the repository at this point in the history
  • Loading branch information
@rhkleijn
rhkleijn committed Jul 2, 2022
1 parent 725258e commit 7d598fb
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 47 deletions.
89 changes: 45 additions & 44 deletions xarray/core/_reductions.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
if TYPE_CHECKING:
from .dataarray import DataArray
from .dataset import Dataset
from .types import T_DataArray, T_Dataset

try:
import flox
Expand All @@ -35,12 +36,12 @@ def reduce(
raise NotImplementedError()

def count(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``count`` along some dimension(s).
Expand Down Expand Up @@ -107,12 +108,12 @@ def count(
)

def all(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``all`` along some dimension(s).
Expand Down Expand Up @@ -179,12 +180,12 @@ def all(
)

def any(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``any`` along some dimension(s).
Expand Down Expand Up @@ -251,13 +252,13 @@ def any(
)

def max(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``max`` along some dimension(s).
Expand Down Expand Up @@ -338,13 +339,13 @@ def max(
)

def min(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``min`` along some dimension(s).
Expand Down Expand Up @@ -425,13 +426,13 @@ def min(
)

def mean(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``mean`` along some dimension(s).
Expand Down Expand Up @@ -516,14 +517,14 @@ def mean(
)

def prod(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
min_count: int | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``prod`` along some dimension(s).
Expand Down Expand Up @@ -623,14 +624,14 @@ def prod(
)

def sum(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
min_count: int | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``sum`` along some dimension(s).
Expand Down Expand Up @@ -730,14 +731,14 @@ def sum(
)

def std(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
ddof: int = 0,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``std`` along some dimension(s).
Expand Down Expand Up @@ -834,14 +835,14 @@ def std(
)

def var(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
ddof: int = 0,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``var`` along some dimension(s).
Expand Down Expand Up @@ -938,13 +939,13 @@ def var(
)

def median(
self,
self: T_Dataset,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> Dataset:
) -> T_Dataset:
"""
Reduce this Dataset's data by applying ``median`` along some dimension(s).
Expand Down Expand Up @@ -1045,12 +1046,12 @@ def reduce(
raise NotImplementedError()

def count(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``count`` along some dimension(s).
Expand Down Expand Up @@ -1111,12 +1112,12 @@ def count(
)

def all(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``all`` along some dimension(s).
Expand Down Expand Up @@ -1177,12 +1178,12 @@ def all(
)

def any(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``any`` along some dimension(s).
Expand Down Expand Up @@ -1243,13 +1244,13 @@ def any(
)

def max(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``max`` along some dimension(s).
Expand Down Expand Up @@ -1322,13 +1323,13 @@ def max(
)

def min(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``min`` along some dimension(s).
Expand Down Expand Up @@ -1401,13 +1402,13 @@ def min(
)

def mean(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``mean`` along some dimension(s).
Expand Down Expand Up @@ -1484,14 +1485,14 @@ def mean(
)

def prod(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
min_count: int | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``prod`` along some dimension(s).
Expand Down Expand Up @@ -1581,14 +1582,14 @@ def prod(
)

def sum(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
min_count: int | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``sum`` along some dimension(s).
Expand Down Expand Up @@ -1678,14 +1679,14 @@ def sum(
)

def std(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
ddof: int = 0,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``std`` along some dimension(s).
Expand Down Expand Up @@ -1772,14 +1773,14 @@ def std(
)

def var(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
ddof: int = 0,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``var`` along some dimension(s).
Expand Down Expand Up @@ -1866,13 +1867,13 @@ def var(
)

def median(
self,
self: T_DataArray,
dim: None | Hashable | Sequence[Hashable] = None,
*,
skipna: bool | None = None,
keep_attrs: bool | None = None,
**kwargs: Any,
) -> DataArray:
) -> T_DataArray:
"""
Reduce this DataArray's data by applying ``median`` along some dimension(s).
Expand Down
Loading

0 comments on commit 7d598fb

Please sign in to comment.