Update expand_args_to_dims

* Generalize type of dim (str -> Hashable)
* Rename to_array -> to_list
* Make error message clearer (only print unexpected inputs)
* Allow lists of length one as inputs
* Add tests

Author: kmsquire

xarray/core/utils.py

@@ -921,10 +921,10 @@ def iterate_nested(nested_list):
 
 
 def expand_args_to_dims(
-    dim: ItemOrSequence[str],
+    dim: ItemOrSequence[Hashable],
     arg_names: Sequence[str],
     args: Sequence[ItemOrSequence[Any]],
-) -> Tuple[Sequence[str], Sequence[Sequence[Any]]]:
+) -> Tuple[Sequence[Hashable], Sequence[Sequence[Any]]]:
     """Expand dims and all elements in args to be arrays of the length of the number of dimensions
 
     Parameters
@@ -946,31 +946,33 @@ def expand_args_to_dims(
     list of dims, list of lists of arguments
     """
     if is_scalar(dim):
-        for name, arg in zip(arg_names, args):
-            if not is_scalar(arg):
-                raise ValueError(f"Expected {name}={arg!r} to be a scalar like 'dim'.")
-        assert isinstance(dim, str) or isinstance(dim, bytes)
-        dim = [dim]
+        dim_list: Sequence[Hashable] = [dim]
+    else:
+        assert isinstance(dim, list)
+        dim_list = dim
 
     # dim is now a list
-    nroll = len(dim)
+    nroll = len(dim_list)
 
-    def to_array(arg):
+    def to_list(arg):
         if is_scalar(arg):
             return [arg] * nroll
+        elif isinstance(arg, list) and len(arg) == 1:
+            return [arg[0]] * nroll
         return arg
 
-    arr_args = [to_array(arg) for arg in args]
+    arr_args = [to_list(arg) for arg in args]
 
-    if any(len(dim) != len(arg) for arg in arr_args):
-        names_vals = ", ".join(
-            f"{name}={val!r}" for name, val in zip(arg_names, arr_args)
+    if any(len(arg) != len(dim_list) for arg in arr_args):
+        names_args_len = ", ".join(
+            f"{name}={args!r} (len={len(args)})" for name, args in zip(arg_names, arr_args)
+            if len(args) != len(dim_list)
         )
         raise ValueError(
-            "Arguments must all be the same length. " f"Received {names_vals}."
+            f"Expected all arguments to have len={len(dim_list)}.  Received: {names_args_len}"
         )
 
-    return dim, arr_args
+    return dim_list, arr_args
 
 
 def get_pads(

xarray/tests/test_utils.py

@@ -8,7 +8,7 @@
 from xarray.coding.cftimeindex import CFTimeIndex
 from xarray.core import duck_array_ops, utils
 from xarray.core.indexes import PandasIndex
-from xarray.core.utils import either_dict_or_kwargs, iterate_nested
+from xarray.core.utils import either_dict_or_kwargs, expand_args_to_dims, iterate_nested
 
 from . import assert_array_equal, requires_cftime, requires_dask
 from .test_coding_times import _all_cftime_date_types
@@ -333,3 +333,24 @@ def test_infix_dims_errors(supplied, all_):
 )
 def test_iterate_nested(nested_list, expected):
     assert list(iterate_nested(nested_list)) == expected
+
+
+def test_expand_args_to_dims():
+    dims, (arg1, arg2, arg3, arg4) = expanded_args = expand_args_to_dims(
+        ["a", "b"],
+        ["arg1", "arg2", "arg3", "arg4"],
+        [1, ["val2.1", "val2.2"], False, [True, False]],
+    )
+
+    assert dims == ["a", "b"]
+    assert arg1 == [1, 1]
+    assert arg2 == ["val2.1", "val2.2"]
+    assert arg3 == [False, False]
+    assert arg4 == [True, False]
+
+    with pytest.raises(ValueError, match="Expected all arguments"):
+        expand_args_to_dims(
+            ["a", "b"],
+            ["arg1", "arg2", "arg3", "arg4"],
+            ["asdf", ["arg2.1", "arg2.2"], ["arg3.1"], ["arg4.1", "arg4.2", "arg4.3"]],
+        )