fixups

Author: jhamman

src/zarr/api/asynchronous.py

@@ -123,7 +123,7 @@ def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple
             # bcolz carray
             chunks = (a.chunklen,) + a.shape[1:]
 
-    return shape, chunks  # type: ignore[return-value]
+    return shape, chunks
 
 
 class _LikeArgs(TypedDict):

src/zarr/core/array.py

@@ -1326,7 +1326,13 @@ def _shard_grid_shape(self) -> tuple[int, ...]:
             The shape of the shard grid for this array.
         """
         if self.shards is None:
-            shard_shape = self.metadata.chunk_grid.chunk_shape
+            chunk_grid = self.metadata.chunk_grid
+            if not isinstance(chunk_grid, RegularChunkGrid):
+                raise NotImplementedError(
+                    "shard_shape is not supported for arrays with variable chunk sizes "
+                    "(RectilinearChunkGrid)."
+                )
+            shard_shape = chunk_grid.chunk_shape
         else:
             shard_shape = self.shards
         return tuple(starmap(ceildiv, zip(self.shape, shard_shape, strict=True)))
@@ -1402,10 +1408,11 @@ async def example():
         if self.shards is None:
             chunks_per_shard = 1
         else:
-            # Sharding only applies to RegularChunkGrid, so chunks is tuple[int, ...]
-            chunks = cast(tuple[int, ...], self.chunks)
+            # Sharding only applies to RegularChunkGrid
+            # Use metadata.chunks to get the inner chunk shape when sharding is used
+            chunk_shape = self.metadata.chunks
             chunks_per_shard = product(
-                tuple(a // b for a, b in zip(self.shards, chunks, strict=True))
+                tuple(a // b for a, b in zip(self.shards, chunk_shape, strict=True))
             )
         return (await self._nshards_initialized()) * chunks_per_shard
 
@@ -5221,7 +5228,15 @@ def _parse_keep_array_attr(
 ]:
     if isinstance(data, Array):
         if chunks == "keep":
-            chunks = data.chunks  # type: ignore[assignment]
+            # Get the chunk shape(s) from the chunk_grid
+            chunk_grid = data.chunk_grid
+            if isinstance(chunk_grid, RegularChunkGrid):
+                chunks = chunk_grid.chunk_shape
+            elif isinstance(chunk_grid, RectilinearChunkGrid):
+                chunks = chunk_grid.chunk_shapes  # type: ignore[assignment]
+            else:
+                msg = f"Unsupported chunk grid type: {type(chunk_grid).__name__}"
+                raise TypeError(msg)
         if shards == "keep":
             shards = data.shards
         if zarr_format is None:
@@ -5701,14 +5716,15 @@ def _iter_shard_regions(
         If the array uses RectilinearChunkGrid (variable-sized chunks).
         when no shards are present.
     """
-    if isinstance(array.metadata.chunk_grid, RectilinearChunkGrid):
+    chunk_grid = array.metadata.chunk_grid
+    if not isinstance(chunk_grid, RegularChunkGrid):
         raise NotImplementedError(
             "_iter_shard_regions is not supported for arrays with variable-sized chunks "
             "(RectilinearChunkGrid). Use the chunk_grid API directly for variable chunk access."
         )
 
     if array.shards is None:
-        shard_shape: Sequence[int] = array.metadata.chunk_grid.chunk_shape
+        shard_shape: Sequence[int] = chunk_grid.chunk_shape
     else:
         shard_shape = array.shards
 
@@ -5747,15 +5763,16 @@ def _iter_chunk_regions(
     NotImplementedError
         If the array uses RectilinearChunkGrid (variable-sized chunks).
     """
-    if isinstance(array.metadata.chunk_grid, RectilinearChunkGrid):
+    chunk_grid = array.metadata.chunk_grid
+    if not isinstance(chunk_grid, RegularChunkGrid):
         raise NotImplementedError(
             "_iter_chunk_regions is not supported for arrays with variable-sized chunks "
             "(RectilinearChunkGrid). Use the chunk_grid API directly for variable chunk access."
         )
 
     return _iter_regions(
         array.shape,
-        array.metadata.chunk_grid.chunk_shape,
+        chunk_grid.chunk_shape,
         origin=origin,
         selection_shape=selection_shape,
         trim_excess=True,

src/zarr/core/chunk_grids.py

@@ -6,7 +6,7 @@
 import numbers
 import operator
 import warnings
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from collections.abc import Sequence
 from dataclasses import dataclass, replace
 from functools import cached_property, reduce

tests/test_api.py

@@ -154,7 +154,7 @@ async def test_array_like_creation(
         kwargs["chunks"] = out_chunks
         expect_chunks = out_chunks
     else:
-        expect_chunks = ref_arr.chunks  # type: ignore[assignment]
+        expect_chunks = ref_arr.chunks
     if out_dtype != "keep":
         kwargs["dtype"] = out_dtype
         expect_dtype = out_dtype

tests/test_chunk_grids/test_rectilinear.py

@@ -9,6 +9,7 @@
 import zarr
 from zarr.core.chunk_grids import (
     RectilinearChunkGrid,
+    RegularChunkGrid,
     _compress_run_length_encoding,
     _expand_run_length_encoding,
     _parse_chunk_shapes,
@@ -1280,7 +1281,9 @@ async def test_array_chunks_property_rectilinear_raises() -> None:
         _ = arr.chunks
 
     # Use chunk_grid instead
-    assert arr.chunk_grid.chunk_shapes == ((10, 20, 30), (25, 25))
+    chunk_grid = arr.chunk_grid
+    assert isinstance(chunk_grid, RectilinearChunkGrid)
+    assert chunk_grid.chunk_shapes == ((10, 20, 30), (25, 25))
 
 
 async def test_array_chunk_grid_property_rectilinear() -> None:
@@ -1318,7 +1321,9 @@ def test_sync_array_chunks_property_rectilinear_raises() -> None:
         _ = arr.chunks
 
     # Use chunk_grid instead
-    assert arr.chunk_grid.chunk_shapes == ((10, 20, 30, 40), (25, 25, 50))
+    chunk_grid = arr.chunk_grid
+    assert isinstance(chunk_grid, RectilinearChunkGrid)
+    assert chunk_grid.chunk_shapes == ((10, 20, 30, 40), (25, 25, 50))
 
 
 async def test_array_chunks_property_regular_with_warning() -> None:
@@ -1349,7 +1354,9 @@ async def test_array_chunks_property_regular_with_warning() -> None:
     assert isinstance(chunks[1], int)
 
     # chunk_grid.chunk_shape works without warning
-    assert arr.chunk_grid.chunk_shape == (10, 20)
+    chunk_grid = arr.chunk_grid
+    assert isinstance(chunk_grid, RegularChunkGrid)
+    assert chunk_grid.chunk_shape == (10, 20)
 
 
 async def test_array_chunk_grid_after_resize() -> None:
@@ -1364,13 +1371,17 @@ async def test_array_chunk_grid_after_resize() -> None:
         zarr_format=3,
     )
 
-    assert arr.chunk_grid.chunk_shapes == ((10, 20), (15, 15))
+    chunk_grid = arr.chunk_grid
+    assert isinstance(chunk_grid, RectilinearChunkGrid)
+    assert chunk_grid.chunk_shapes == ((10, 20), (15, 15))
 
     # Resize to grow
     await arr.resize((50, 45))
 
     # Chunk grid should be updated
-    assert arr.chunk_grid.chunk_shapes == ((10, 20, 20), (15, 15, 15))
+    chunk_grid = arr.chunk_grid
+    assert isinstance(chunk_grid, RectilinearChunkGrid)
+    assert chunk_grid.chunk_shapes == ((10, 20, 20), (15, 15, 15))
 
 
 async def test_metadata_chunks_property_raises_for_rectilinear() -> None:
@@ -1390,7 +1401,9 @@ async def test_metadata_chunks_property_raises_for_rectilinear() -> None:
         _ = arr.metadata.chunks
 
     # Use chunk_grid instead
-    assert arr.metadata.chunk_grid.chunk_shapes == ((10, 20, 30), (25, 25))
+    chunk_grid = arr.metadata.chunk_grid
+    assert isinstance(chunk_grid, RectilinearChunkGrid)
+    assert chunk_grid.chunk_shapes == ((10, 20, 30), (25, 25))
 
 
 # ===================================================================