add parse_dims_as_set

Author: shoyer

xarray/core/computation.py

@@ -31,7 +31,7 @@
 from xarray.core.merge import merge_attrs, merge_coordinates_without_align
 from xarray.core.options import OPTIONS, _get_keep_attrs
 from xarray.core.types import Dims, T_DataArray
-from xarray.core.utils import is_dict_like, is_scalar, parse_dims
+from xarray.core.utils import is_dict_like, is_scalar, parse_dims_as_set
 from xarray.core.variable import Variable
 from xarray.namedarray.parallelcompat import get_chunked_array_type
 from xarray.namedarray.pycompat import is_chunked_array
@@ -1846,11 +1846,9 @@ def dot(
         dim_counts: Counter = Counter()
         for arr in arrays:
             dim_counts.update(arr.dims)
-        dim = tuple(d for d, c in dim_counts.items() if c > 1)
+        dot_dims = {d for d, c in dim_counts.items() if c > 1}
     else:
-        dim = parse_dims(dim, all_dims=tuple(all_dims))
-
-    dot_dims: set[Hashable] = set(dim)
+        dot_dims = parse_dims_as_set(dim, all_dims=set(all_dims))
 
     # dimensions to be parallelized
     broadcast_dims = common_dims - dot_dims

xarray/core/dataset.py

@@ -118,7 +118,7 @@
     is_duck_dask_array,
     is_scalar,
     maybe_wrap_array,
-    parse_dims,
+    parse_dims_as_set,
 )
 from xarray.core.variable import (
     IndexVariable,
@@ -6987,7 +6987,7 @@ def reduce(
                 " Please use 'dim' instead."
             )
 
-        dims = set(parse_dims(dim, tuple(self.dims)))
+        dims = parse_dims_as_set(dim, self._dims.keys())
 
         if keep_attrs is None:
             keep_attrs = _get_keep_attrs(default=False)

xarray/core/datatree.py

@@ -43,7 +43,7 @@
     drop_dims_from_indexers,
     either_dict_or_kwargs,
     maybe_wrap_array,
-    parse_dims,
+    parse_dims_as_set,
 )
 from xarray.core.variable import Variable
 
@@ -1631,7 +1631,7 @@ def reduce(
         **kwargs: Any,
     ) -> Self:
         """Reduce this tree by applying `func` along some dimension(s)."""
-        dims = set(parse_dims(dim, tuple(self._get_all_dims())))
+        dims = parse_dims_as_set(dim, self._get_all_dims())
         result = {}
         for node in self.subtree:
             reduce_dims = [d for d in node._node_dims if d in dims]

xarray/core/utils.py

@@ -58,6 +58,7 @@
     Mapping,
     MutableMapping,
     MutableSet,
+    Set,
     Sequence,
     ValuesView,
 )
@@ -831,7 +832,7 @@ def drop_dims_from_indexers(
 
 
 @overload
-def parse_dims(
+def parse_dims_as_tuple(
     dim: Dims,
     all_dims: tuple[Hashable, ...],
     *,
@@ -841,7 +842,7 @@ def parse_dims(
 
 
 @overload
-def parse_dims(
+def parse_dims_as_tuple(
     dim: Dims,
     all_dims: tuple[Hashable, ...],
     *,
@@ -850,7 +851,7 @@ def parse_dims(
 ) -> tuple[Hashable, ...] | None | EllipsisType: ...
 
 
-def parse_dims(
+def parse_dims_as_tuple(
     dim: Dims,
     all_dims: tuple[Hashable, ...],
     *,
@@ -891,6 +892,47 @@ def parse_dims(
     return tuple(dim)
 
 
+@overload
+def parse_dims_as_set(
+    dim: Dims,
+    all_dims: Set[Hashable],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[True] = True,
+) -> Set[Hashable]: ...
+
+
+@overload
+def parse_dims_as_set(
+    dim: Dims,
+    all_dims: Set[Hashable],
+    *,
+    check_exists: bool = True,
+    replace_none: Literal[False],
+) -> Set[Hashable] | None | EllipsisType: ...
+
+
+def parse_dims_as_set(
+    dim: Dims,
+    all_dims: Set[Hashable],
+    *,
+    check_exists: bool = True,
+    replace_none: bool = True,
+) -> Set[Hashable] | None | EllipsisType:
+    """Like parse_dims_as_tuple, but returning a set instead of a tuple."""
+    # TODO: Consider removing parse_dims_as_tuple?
+    if dim is None or dim is ...:
+        if replace_none:
+            return all_dims
+        return dim
+    if isinstance(dim, str):
+        dim = {dim}
+    dim = set(dim)
+    if check_exists:
+        _check_dims(dim, all_dims)
+    return dim
+
+
 @overload
 def parse_ordered_dims(
     dim: Dims,
@@ -958,15 +1000,15 @@ def parse_ordered_dims(
         return dims[:idx] + other_dims + dims[idx + 1 :]
     else:
         # mypy cannot resolve that the sequence cannot contain "..."
-        return parse_dims(  # type: ignore[call-overload]
+        return parse_dims_as_tuple(  # type: ignore[call-overload]
             dim=dim,
             all_dims=all_dims,
             check_exists=check_exists,
             replace_none=replace_none,
         )
 
 
-def _check_dims(dim: set[Hashable], all_dims: set[Hashable]) -> None:
+def _check_dims(dim: Set[Hashable], all_dims: Set[Hashable]) -> None:
     wrong_dims = (dim - all_dims) - {...}
     if wrong_dims:
         wrong_dims_str = ", ".join(f"'{d!s}'" for d in wrong_dims)

xarray/tests/test_utils.py

@@ -283,16 +283,16 @@ def test_infix_dims_errors(supplied, all_):
         pytest.param(..., ..., id="ellipsis"),
     ],
 )
-def test_parse_dims(dim, expected) -> None:
+def test_parse_dims_as_tuple(dim, expected) -> None:
     all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
-    actual = utils.parse_dims(dim, all_dims, replace_none=False)
+    actual = utils.parse_dims_as_tuple(dim, all_dims, replace_none=False)
     assert actual == expected
 
 
 def test_parse_dims_set() -> None:
     all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
     dim = {"a", 1}
-    actual = utils.parse_dims(dim, all_dims)
+    actual = utils.parse_dims_as_tuple(dim, all_dims)
     assert set(actual) == dim
 
 
@@ -301,7 +301,7 @@ def test_parse_dims_set() -> None:
 )
 def test_parse_dims_replace_none(dim: None | EllipsisType) -> None:
     all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
-    actual = utils.parse_dims(dim, all_dims, replace_none=True)
+    actual = utils.parse_dims_as_tuple(dim, all_dims, replace_none=True)
     assert actual == all_dims
 
 
@@ -316,7 +316,7 @@ def test_parse_dims_replace_none(dim: None | EllipsisType) -> None:
 def test_parse_dims_raises(dim) -> None:
     all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
     with pytest.raises(ValueError, match="'x'"):
-        utils.parse_dims(dim, all_dims, check_exists=True)
+        utils.parse_dims_as_tuple(dim, all_dims, check_exists=True)
 
 
 @pytest.mark.parametrize(