Support a single chunk spec in DataTree.chunk

doc/whats-new.rst

@@ -117,6 +117,11 @@ Bug Fixes
   type error when applying reduction methods, due to the reduction methods being
   dynamically generated (:issue:`8136`).
   By `Andrew Scherer <https://github.com/andrew-s28>`_.
+- :py:meth:`DataTree.chunk` now accepts a single chunk specification such as
+  ``"auto"`` or an integer and applies it to every dimension in the tree, as
+  :py:meth:`Dataset.chunk` does, instead of raising ``TypeError``
+  (:issue:`11315`, :pull:`11569`).
+  By `fredrikblau <https://github.com/fredrikblau>`_.
 
 .. _`pandas-dev/pandas#64793`: https://github.com/pandas-dev/pandas/pull/64793
 

xarray/core/datatree.py

@@ -2613,9 +2613,10 @@ def chunk(
 
         Parameters
         ----------
-        chunks : int, tuple of int, "auto" or mapping of hashable to int or a TimeResampler, optional
+        chunks : int, "auto" or mapping of hashable to int or a TimeResampler, optional
             Chunk sizes along each dimension, e.g., ``5``, ``"auto"``, or
             ``{"x": 5, "y": 5}`` or ``{"x": 5, "time": TimeResampler(freq="YE")}``.
+            A single value is applied to every dimension in the tree.
         name_prefix : str, default: "xarray-"
             Prefix for the name of any new dask arrays.
         token : str, optional
@@ -2650,15 +2651,24 @@ def chunk(
         xarray.unify_chunks
         dask.array.from_array
         """
-        # don't support deprecated ways of passing chunks
-        if not isinstance(chunks, Mapping):
-            raise TypeError(
-                f"invalid type for chunks: {type(chunks)}. Only mappings are supported."
-            )
-        combined_chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
-
         all_dims = self._get_all_dims()
 
+        combined_chunks: Mapping[Any, T_ChunkDimFreq]
+        if not isinstance(chunks, Mapping):
+            # don't support deprecated ways of passing chunks: sequences of
+            # dimension-order sizes are deprecated for Dataset.chunk, and are
+            # additionally ambiguous for a tree whose groups need not share an
+            # ordering of their dimensions
+            if chunks is None or isinstance(chunks, tuple | list):
+                raise TypeError(
+                    f"invalid type for chunks: {type(chunks)}. Only mappings and "
+                    'single chunk specifications (e.g. 5 or "auto") to be applied '
+                    "to every dimension are supported."
+                )
+            combined_chunks = dict.fromkeys(all_dims, chunks)
+        else:
+            combined_chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
+
         bad_dims = combined_chunks.keys() - all_dims
         if bad_dims:
             raise ValueError(

xarray/tests/test_datatree.py

@@ -2784,3 +2784,20 @@ def test_chunk(self):
 
         with pytest.raises(ValueError, match="not found in data dimensions"):
             tree.chunk({"u": 2})
+
+    @pytest.mark.parametrize("chunks", ["auto", 5, "20B"])
+    def test_chunk_single_spec(self, chunks):
+        ds1 = xr.Dataset({"a": ("x", np.arange(10))})
+        ds2 = xr.Dataset({"b": ("y", np.arange(6))})
+
+        tree = xr.DataTree.from_dict({"/": ds1, "/group1": ds2})
+        actual = tree.chunk(chunks)
+
+        expected = xr.DataTree.from_dict(
+            {"/": ds1.chunk(chunks), "/group1": ds2.chunk(chunks)}
+        )
+
+        assert_identical(actual, expected)
+        assert actual.chunksizes == expected.chunksizes
+        assert set(actual.chunksizes["/"]) == {"x"}
+        assert set(actual.chunksizes["/group1"]) == {"y"}