fix: validate that dask-style chunks have regular shapes (#3779)

* fix: validate that dask-style chunks have regular shapes

* Apply suggestion from @dcherian

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

* Fix shortcircuit option

---------

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>

Author: maxrjones

src/zarr/core/chunk_grids.py

@@ -126,11 +126,14 @@ def normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tupl
         chunks = tuple(int(chunks) for _ in shape)
 
     # handle dask-style chunks (iterable of iterables)
-    if all(isinstance(c, (tuple | list)) for c in chunks):
-        # take first chunk size for each dimension
-        chunks = tuple(
-            c[0] for c in chunks
-        )  # TODO: check/error/warn for irregular chunks (e.g. if c[0] != c[1:-1])
+    if all(isinstance(c, (tuple, list)) for c in chunks):
+        for i, c in enumerate(chunks):
+            if any(x != y for x, y in itertools.pairwise(c[:-1])) or (len(c) > 1 and c[-1] > c[0]):
+                raise ValueError(
+                    f"Irregular chunk sizes in dimension {i}: {tuple(c)}. "
+                    "Only uniform chunks (with an optional smaller final chunk) are supported."
+                )
+        chunks = tuple(c[0] for c in chunks)
 
     # handle bad dimensionality
     if len(chunks) > len(shape):

tests/test_chunk_grids.py

@@ -35,6 +35,10 @@ def test_guess_chunks(shape: tuple[int, ...], itemsize: int) -> None:
         ((30, None, None), (100, 20, 10), 1, (30, 20, 10)),
         ((30, 20, None), (100, 20, 10), 1, (30, 20, 10)),
         ((30, 20, 10), (100, 20, 10), 1, (30, 20, 10)),
+        # dask-style chunks (uniform with optional smaller final chunk)
+        (((100, 100, 100), (50, 50)), (300, 100), 1, (100, 50)),
+        (((100, 100, 50),), (250,), 1, (100,)),
+        (((100,),), (100,), 1, (100,)),
         # auto chunking
         (None, (100,), 1, (100,)),
         (-1, (100,), 1, (100,)),
@@ -52,3 +56,8 @@ def test_normalize_chunks_errors() -> None:
         normalize_chunks("foo", (100,), 1)
     with pytest.raises(ValueError):
         normalize_chunks((100, 10), (100,), 1)
+    # dask-style irregular chunks should raise
+    with pytest.raises(ValueError, match="Irregular chunk sizes"):
+        normalize_chunks(((10, 20, 30),), (60,), 1)
+    with pytest.raises(ValueError, match="Irregular chunk sizes"):
+        normalize_chunks(((100, 100), (10, 20)), (200, 30), 1)