Support a single chunk spec in DataTree.chunk

[fredrikblau]

• Closes Cannot call .chunk('auto') on DataTree. #11315
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst

DataTree.chunk rejects anything that isn't a mapping:

>>> ds = xr.Dataset({"a": ("x", np.arange(10))})
>>> ds.chunk("auto").chunksizes
Frozen({'x': (10,)})
>>> xr.DataTree(ds).chunk("auto")
TypeError: invalid type for chunks: <class 'str'>. Only mappings are supported.

That contradicts the method's own docstring and its chunks: T_ChunksFreq annotation, both of which advertise int, "auto" or a mapping, and it diverges from Dataset.chunk, which broadcasts a single value across every dimension with dict.fromkeys(self.dims, chunks).

This does the same for trees: a non-mapping chunks is expanded with dict.fromkeys(self._get_all_dims(), chunks) before the existing per-group dispatch, so dt.chunk("auto"), dt.chunk(5) and dt.chunk("20B") work and each group gets chunk sizes for its own dimensions only.

The existing guard was commented "don't support deprecated ways of passing chunks", and that part is kept. Tuples and lists still raise TypeError: they are the form Dataset.chunk emits a FutureWarning for, and for a tree they are ambiguous anyway, since groups need not share an ordering of their dimensions and there is no well-defined sequence to zip against _get_all_dims(). None also still raises, matching the behaviour test_chunk already asserts. The message is reworded to say what is accepted while keeping the invalid type for chunks: prefix, so both existing pytest.raises(TypeError, match="invalid type") assertions still pass unchanged.

tuple of int was dropped from the docstring's chunks line, since that form is deliberately not accepted here.

I left T_ChunksFreq alone — it still nominally includes tuple and None, but it is shared with Dataset.chunk and narrowing it would be a wider change than this fix warrants.

Testing

TestDask::test_chunk_single_spec is parametrized over "auto", 5 and "20B" and compares against the per-group Dataset.chunk result rather than just asserting no exception, so it pins the actual chunk sizes: tree.chunk(5) gives {'/': {'x': (5, 5)}, '/group1': {'y': (5, 1)}} and tree.chunk("20B") gives {'/': {'x': (2, 2, 2, 2, 2)}, '/group1': {'y': (2, 2, 2)}} — different per group, which is what a broadcast should produce.

Reverting datatree.py alone fails all three parametrizations with the original TypeError; with the change they pass.

xarray/tests/test_datatree.py: 159 passed, 5 xfailed. xarray/tests/test_dask.py: 219 passed, 2 skipped, 3 xfailed.