diff --git a/xarray/core/types.py b/xarray/core/types.py index 41078d29c0e..5ec76046e8a 100644 --- a/xarray/core/types.py +++ b/xarray/core/types.py @@ -182,7 +182,8 @@ def copy( Dims = Union[str, Collection[Hashable], "ellipsis", None] # FYI in some cases we don't allow `None`, which this doesn't take account of. -T_ChunkDim: TypeAlias = Union[int, Literal["auto"], None, tuple[int, ...]] +# FYI the `str` is for a size string, e.g. "16MB", supported by dask. +T_ChunkDim: TypeAlias = Union[str, int, Literal["auto"], None, tuple[int, ...]] # We allow the tuple form of this (though arguably we could transition to named dims only) T_Chunks: TypeAlias = Union[T_ChunkDim, Mapping[Any, T_ChunkDim]] T_NormalizedChunks = tuple[tuple[int, ...], ...] diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py index 517fc0c2d62..3ac638c3c5f 100644 --- a/xarray/tests/test_dask.py +++ b/xarray/tests/test_dask.py @@ -337,13 +337,16 @@ def setUp(self): self.data, coords={"x": range(4)}, dims=("x", "y"), name="foo" ) - def test_chunk(self): + def test_chunk(self) -> None: for chunks, expected in [ ({}, ((2, 2), (2, 2, 2))), (3, ((3, 1), (3, 3))), ({"x": 3, "y": 3}, ((3, 1), (3, 3))), ({"x": 3}, ((3, 1), (2, 2, 2))), ({"x": (3, 1)}, ((3, 1), (2, 2, 2))), + ({"x": "16B"}, ((1, 1, 1, 1), (2, 2, 2))), + ("16B", ((1, 1, 1, 1), (1,) * 6)), + ("16MB", ((4,), (6,))), ]: # Test DataArray rechunked = self.lazy_array.chunk(chunks)