Add parse_dims func (#7051)

* add parse_dims func * add some more tests for tuples * add parse_ordered_dims * fix typing issue * remove double ellipsis typehints * fix gen aggrs * remove more double ellipsis typehints * fix doctests: supress urllib3 warning (#7326) * Enable `origin` and `offset` arguments in `resample` (#7284) * Initial work toward enabling origin and offset arguments in resample * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix _convert_offset_to_timedelta * Reduce number of tests * Address initial review comments * Add more typing information * Make cftime import lazy * Fix module_available import and test * Remove old origin argument * Add type annotations for resample_cftime.py * Add None as a possibility for closed and label * Add what's new entry * Add missing type annotation * Delete added line * Fix typing errors * Add comment and test for as_timedelta stub * Remove old code * [test-upstream] Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Deepak Cherian <[email protected]> * Fix PR number in what’s new (#7331) * [pre-commit.ci] pre-commit autoupdate (#7330) updates: - [github.com/pre-commit/pre-commit-hooks: v4.3.0 → v4.4.0](pre-commit/pre-commit-hooks@v4.3.0...v4.4.0) - [github.com/PyCQA/autoflake: v1.7.7 → v2.0.0](PyCQA/autoflake@v1.7.7...v2.0.0) - [github.com/PyCQA/flake8: 5.0.4 → 6.0.0](PyCQA/flake8@5.0.4...6.0.0) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Mathias Hauser <[email protected]> * add comment explaining type: ignore * fix doctest win/linux issue once again Co-authored-by: Joe Hamman <[email protected]> Co-authored-by: Mathias Hauser <[email protected]> Co-authored-by: Spencer Clark <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Deepak Cherian <[email protected]>
pydata · Nov 30, 2022 · e4fe194 · e4fe194
1 parent 675a3ff
commit e4fe194
Show file tree

Hide file tree

Showing 11 changed files with 383 additions and 142 deletions.
diff --git a/xarray/core/_aggregations.py b/xarray/core/_aggregations.py
diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py
@@ -3614,7 +3614,7 @@ def combine_first(self: T_DataArray, other: T_DataArray) -> T_DataArray:
     def reduce(
         self: T_DataArray,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,
@@ -4601,7 +4601,7 @@ def imag(self: T_DataArray) -> T_DataArray:
     def dot(
         self: T_DataArray,
         other: T_DataArray,
-        dims: Dims | ellipsis = None,
+        dims: Dims = None,
     ) -> T_DataArray:
         """Perform dot product of two DataArrays along their shared dims.
 
@@ -5605,7 +5605,7 @@ def idxmax(
     # https://github.com/python/mypy/issues/12846 is resolved
     def argmin(
         self,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         axis: int | None = None,
         keep_attrs: bool | None = None,
         skipna: bool | None = None,
@@ -5707,7 +5707,7 @@ def argmin(
     # https://github.com/python/mypy/issues/12846 is resolved
     def argmax(
         self,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         axis: int | None = None,
         keep_attrs: bool | None = None,
         skipna: bool | None = None,

diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py
@@ -5798,7 +5798,7 @@ def combine_first(self: T_Dataset, other: T_Dataset) -> T_Dataset:
     def reduce(
         self: T_Dataset,
         func: Callable,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         keep_attrs: bool | None = None,
         keepdims: bool = False,

diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py
@@ -493,7 +493,7 @@ def map(
     def reduce(
         self,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,
@@ -652,7 +652,7 @@ def _maybe_unstack(self, obj):
 
     def _flox_reduce(
         self,
-        dim: Dims | ellipsis,
+        dim: Dims,
         keep_attrs: bool | None = None,
         **kwargs: Any,
     ):
@@ -1143,7 +1143,7 @@ def _combine(self, applied, shortcut=False):
     def reduce(
         self,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,
@@ -1296,7 +1296,7 @@ def _combine(self, applied):
     def reduce(
         self,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,

diff --git a/xarray/core/resample.py b/xarray/core/resample.py
@@ -49,7 +49,7 @@ def __init__(
 
     def _flox_reduce(
         self,
-        dim: Dims | ellipsis,
+        dim: Dims,
         keep_attrs: bool | None = None,
         **kwargs,
     ) -> T_Xarray:
@@ -368,7 +368,7 @@ def apply(self, func, args=(), shortcut=None, **kwargs):
     def reduce(
         self,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         *,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,

diff --git a/xarray/core/types.py b/xarray/core/types.py
@@ -113,7 +113,8 @@ def dtype(self) -> np.dtype:
 VarCompatible = Union["Variable", "ScalarOrArray"]
 GroupByIncompatible = Union["Variable", "GroupBy"]
 
-Dims = Union[str, Iterable[Hashable], None]
+Dims = Union[str, Iterable[Hashable], "ellipsis", None]
+OrderedDims = Union[str, Sequence[Union[Hashable, "ellipsis"]], "ellipsis", None]
 
 ErrorOptions = Literal["raise", "ignore"]
 ErrorOptionsWithWarn = Literal["raise", "warn", "ignore"]

diff --git a/xarray/core/utils.py b/xarray/core/utils.py
@@ -57,9 +57,11 @@
     Hashable,
     Iterable,
     Iterator,
+    Literal,
     Mapping,
     MutableMapping,
     MutableSet,
+    Sequence,
     TypeVar,
     cast,
     overload,
@@ -69,7 +71,7 @@
 import pandas as pd
 
 if TYPE_CHECKING:
-    from .types import ErrorOptionsWithWarn
+    from .types import Dims, ErrorOptionsWithWarn, OrderedDims
 
 K = TypeVar("K")
 V = TypeVar("V")
@@ -887,15 +889,17 @@ def drop_dims_from_indexers(
 
 
 def drop_missing_dims(
-    supplied_dims: Collection, dims: Collection, missing_dims: ErrorOptionsWithWarn
-) -> Collection:
+    supplied_dims: Iterable[Hashable],
+    dims: Iterable[Hashable],
+    missing_dims: ErrorOptionsWithWarn,
+) -> Iterable[Hashable]:
     """Depending on the setting of missing_dims, drop any dimensions from supplied_dims that
     are not present in dims.
 
     Parameters
     ----------
-    supplied_dims : dict
-    dims : sequence
+    supplied_dims : Iterable of Hashable
+    dims : Iterable of Hashable
     missing_dims : {"raise", "warn", "ignore"}
     """
 
@@ -928,6 +932,158 @@ def drop_missing_dims(
         )
 
 
+T_None = TypeVar("T_None", None, "ellipsis")
+
+
+@overload
+def parse_dims(
+    dim: str | Iterable[Hashable] | T_None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[True] = True,
+) -> tuple[Hashable, ...]:
+    ...
+
+
+@overload
+def parse_dims(
+    dim: str | Iterable[Hashable] | T_None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[False],
+) -> tuple[Hashable, ...] | T_None:
+    ...
+
+
+def parse_dims(
+    dim: Dims,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: bool = True,
+) -> tuple[Hashable, ...] | None | ellipsis:
+    """Parse one or more dimensions.
+
+    A single dimension must be always a str, multiple dimensions
+    can be Hashables. This supports e.g. using a tuple as a dimension.
+    If you supply e.g. a set of dimensions the order cannot be
+    conserved, but for sequences it will be.
+
+    Parameters
+    ----------
+    dim : str, Iterable of Hashable, "..." or None
+        Dimension(s) to parse.
+    all_dims : tuple of Hashable
+        All possible dimensions.
+    check_exists: bool, default: True
+        if True, check if dim is a subset of all_dims.
+    replace_none : bool, default: True
+        If True, return all_dims if dim is None or "...".
+
+    Returns
+    -------
+    parsed_dims : tuple of Hashable
+        Input dimensions as a tuple.
+    """
+    if dim is None or dim is ...:
+        if replace_none:
+            return all_dims
+        return dim
+    if isinstance(dim, str):
+        dim = (dim,)
+    if check_exists:
+        _check_dims(set(dim), set(all_dims))
+    return tuple(dim)
+
+
+@overload
+def parse_ordered_dims(
+    dim: str | Sequence[Hashable | ellipsis] | T_None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[True] = True,
+) -> tuple[Hashable, ...]:
+    ...
+
+
+@overload
+def parse_ordered_dims(
+    dim: str | Sequence[Hashable | ellipsis] | T_None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[False],
+) -> tuple[Hashable, ...] | T_None:
+    ...
+
+
+def parse_ordered_dims(
+    dim: OrderedDims,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check_exists: bool = True,
+    replace_none: bool = True,
+) -> tuple[Hashable, ...] | None | ellipsis:
+    """Parse one or more dimensions.
+
+    A single dimension must be always a str, multiple dimensions
+    can be Hashables. This supports e.g. using a tuple as a dimension.
+    An ellipsis ("...") in a sequence of dimensions will be
+    replaced with all remaining dimensions. This only makes sense when
+    the input is a sequence and not e.g. a set.
+
+    Parameters
+    ----------
+    dim : str, Sequence of Hashable or "...", "..." or None
+        Dimension(s) to parse. If "..." appears in a Sequence
+        it always gets replaced with all remaining dims
+    all_dims : tuple of Hashable
+        All possible dimensions.
+    check_exists: bool, default: True
+        if True, check if dim is a subset of all_dims.
+    replace_none : bool, default: True
+        If True, return all_dims if dim is None.
+
+    Returns
+    -------
+    parsed_dims : tuple of Hashable
+        Input dimensions as a tuple.
+    """
+    if dim is not None and dim is not ... and not isinstance(dim, str) and ... in dim:
+        dims_set: set[Hashable | ellipsis] = set(dim)
+        all_dims_set = set(all_dims)
+        if check_exists:
+            _check_dims(dims_set, all_dims_set)
+        if len(all_dims_set) != len(all_dims):
+            raise ValueError("Cannot use ellipsis with repeated dims")
+        dims = tuple(dim)
+        if dims.count(...) > 1:
+            raise ValueError("More than one ellipsis supplied")
+        other_dims = tuple(d for d in all_dims if d not in dims_set)
+        idx = dims.index(...)
+        return dims[:idx] + other_dims + dims[idx + 1 :]
+    else:
+        # mypy cannot resolve that the sequence cannot contain "..."
+        return parse_dims(  # type: ignore[call-overload]
+            dim=dim,
+            all_dims=all_dims,
+            check_exists=check_exists,
+            replace_none=replace_none,
+        )
+
+
+def _check_dims(dim: set[Hashable | ellipsis], all_dims: set[Hashable]) -> None:
+    wrong_dims = dim - all_dims
+    if wrong_dims and wrong_dims != {...}:
+        wrong_dims_str = ", ".join(f"'{d!s}'" for d in wrong_dims)
+        raise ValueError(
+            f"Dimension(s) {wrong_dims_str} do not exist. Expected one or more of {all_dims}"
+        )
+
+
 _Accessor = TypeVar("_Accessor")
 
 

diff --git a/xarray/core/variable.py b/xarray/core/variable.py
@@ -1889,7 +1889,7 @@ def clip(self, min=None, max=None):
     def reduce(
         self,
         func: Callable[..., Any],
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         axis: int | Sequence[int] | None = None,
         keep_attrs: bool | None = None,
         keepdims: bool = False,
@@ -2663,7 +2663,7 @@ def _to_numeric(self, offset=None, datetime_unit=None, dtype=float):
     def _unravel_argminmax(
         self,
         argminmax: str,
-        dim: Dims | ellipsis,
+        dim: Dims,
         axis: int | None,
         keep_attrs: bool | None,
         skipna: bool | None,
@@ -2732,7 +2732,7 @@ def _unravel_argminmax(
 
     def argmin(
         self,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         axis: int | None = None,
         keep_attrs: bool | None = None,
         skipna: bool | None = None,
@@ -2777,7 +2777,7 @@ def argmin(
 
     def argmax(
         self,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         axis: int | None = None,
         keep_attrs: bool | None = None,
         skipna: bool | None = None,

diff --git a/xarray/core/weighted.py b/xarray/core/weighted.py
@@ -207,7 +207,7 @@ def _check_dim(self, dim: Dims):
     def _reduce(
         da: DataArray,
         weights: DataArray,
-        dim: Dims | ellipsis = None,
+        dim: Dims = None,
         skipna: bool | None = None,
     ) -> DataArray:
         """reduce using dot; equivalent to (da * weights).sum(dim, skipna)