Add parse_dims func

xarray/core/utils.py

@@ -21,6 +21,7 @@
     Hashable,
     Iterable,
     Iterator,
+    Literal,
     Mapping,
     MutableMapping,
     MutableSet,
@@ -919,6 +920,72 @@ def drop_missing_dims(
         )
 
 
+@overload
+def parse_dims(
+    dim: str | Iterable[Hashable] | None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check: bool = True,
+    replace_none: Literal[True] = True,
+) -> tuple[Hashable, ...]:
+    ...
+
+
+@overload
+def parse_dims(
+    dim: str | Iterable[Hashable] | None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check: bool = True,
+    replace_none: Literal[False],
+) -> tuple[Hashable, ...] | None:
+    ...
+
+
+def parse_dims(
+    dim: str | Iterable[Hashable] | None,
+    all_dims: tuple[Hashable, ...],
+    *,
+    check: bool = True,
+    replace_none: bool = True,
+) -> tuple[Hashable, ...] | None:
+    """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.
+
+    Parameters
+    ----------
+    dim : str, Iterable of Hashable or None
+        Dimension(s) to parse.
+    all_dims : tuple of Hashable
+        All possible dimensions.
+    check: 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 None:
+        if replace_none:
+            return all_dims
+        return None
+    if isinstance(dim, str):
+        dim = (dim,)
+    if check:
+        wrong_dims = set(dim) - set(all_dims)
+        if 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}"
+            )
+    return tuple(dim)
+
+
 _Accessor = TypeVar("_Accessor")
 
 

xarray/tests/test_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Hashable
+from typing import Hashable, Iterable
 
 import numpy as np
 import pandas as pd
@@ -294,6 +294,47 @@ def test_infix_dims_errors(supplied, all_):
         list(utils.infix_dims(supplied, all_))
 
 
+@pytest.mark.parametrize(
+    ["dim", "expected"],
+    [
+        pytest.param("a", ("a",), id="str"),
+        pytest.param(["a", "b"], ("a", "b"), id="list_of_str"),
+        pytest.param(["a", 1], ("a", 1), id="list_mixed"),
+        pytest.param(("a", "b"), ("a", "b"), id="tuple_of_str"),
+        pytest.param(["a", ("b", "c")], ("a", ("b", "c")), id="list_with_tuple"),
+        pytest.param((("b", "c"),), (("b", "c"),), id="tuple_of_tuple"),
+        pytest.param(None, None, id="None"),
+    ],
+)
+def test_parse_dims(
+    dim: str | Iterable[Hashable] | None,
+    expected: tuple[Hashable, ...],
+) -> None:
+    all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
+    actual = utils.parse_dims(dim, all_dims, replace_none=False)
+    assert actual == expected
+
+
+def test_parse_dims_replace_none() -> None:
+    all_dims = ("a", "b", 1, ("b", "c"))  # selection of different Hashables
+    actual = utils.parse_dims(None, all_dims, replace_none=True)
+    assert actual == all_dims
+
+
+@pytest.mark.parametrize(
+    "dim",
+    [
+        pytest.param("x", id="str_missing"),
+        pytest.param(["a", "x"], id="list_missing_one"),
+        pytest.param(["x", 2], id="list_missing_all"),
+    ],
+)
+def test_parse_dims_raises(dim: str | Iterable[Hashable]) -> 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=True)
+
+
 @pytest.mark.parametrize(
     "nested_list, expected",
     [