fixups

Author: jhamman

src/zarr/core/array.py

@@ -1262,12 +1262,24 @@ def _chunk_grid_shape(self) -> tuple[int, ...]:
         """
         The shape of the chunk grid for this array.
 
+        For arrays with sharding, this returns the grid of inner chunks, not shards.
+        For arrays with RectilinearChunkGrid, this returns the grid shape.
+
         Returns
         -------
         tuple[int, ...]
             The shape of the chunk grid for this array.
         """
-        return tuple(starmap(ceildiv, zip(self.shape, self.chunks, strict=True)))
+        chunks = self.chunks
+        # Handle 0-dimensional arrays
+        if len(chunks) == 0:
+            return ()
+        # For RegularChunkGrid (or sharded with RegularChunkGrid), chunks is tuple[int, ...]
+        if isinstance(chunks[0], int):
+            return tuple(starmap(ceildiv, zip(self.shape, chunks, strict=True)))
+        # For RectilinearChunkGrid, chunks is tuple[tuple[int, ...], ...]
+        # Use the chunk_grid method
+        return self.metadata.chunk_grid.get_chunk_grid_shape(self.shape)
 
     @property
     def _shard_grid_shape(self) -> tuple[int, ...]:
@@ -5612,7 +5624,7 @@ def _iter_shard_regions(
         If the array uses RectilinearChunkGrid (variable-sized chunks).
     """
     chunks = array.chunks
-    if not isinstance(chunks[0], int):
+    if chunks and not isinstance(chunks[0], int):
         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."
@@ -5661,7 +5673,7 @@ def _iter_chunk_regions(
         If the array uses RectilinearChunkGrid (variable-sized chunks).
     """
     chunks = array.chunks
-    if not isinstance(chunks[0], int):
+    if chunks and not isinstance(chunks[0], int):
         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."

src/zarr/core/chunk_grids.py

@@ -300,7 +300,8 @@ def _normalize_chunks(chunks: Any, shape: tuple[int, ...], typesize: int) -> tup
         chunks = tuple(int(chunks) for _ in shape)
 
     # handle dask-style chunks (iterable of iterables)
-    # TODO
+    # Note: Only regular chunks are supported here. Irregular chunks will trigger a warning
+    # and use only the first chunk size per dimension. For true variable chunks, use RectilinearChunkGrid.
     if all(isinstance(c, (tuple | list)) for c in chunks):
         # Check for irregular chunks and warn user
         for dim_idx, c in enumerate(chunks):
@@ -346,7 +347,7 @@ def from_dict(cls, data: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any]) ->
 
         # After isinstance check, data must be dict[str, JSON]
         # Cast needed for older mypy versions that don't narrow types properly
-        data_dict = cast(dict[str, JSON], data)  # type: ignore[redundant-cast]
+        data_dict = cast(dict[str, JSON], data)
         name_parsed, _ = parse_named_configuration(data_dict)
         if name_parsed == "regular":
             return RegularChunkGrid._from_dict(data_dict)
@@ -662,7 +663,46 @@ def to_dict(self) -> dict[str, JSON]:
         }
 
     def update_shape(self, new_shape: tuple[int, ...]) -> Self:
-        """TODO - write docstring"""
+        """
+        Update the RectilinearChunkGrid to accommodate a new array shape.
+
+        When resizing an array, this method adjusts the chunk grid to match the new shape.
+        For dimensions that grow, a new chunk is added with size equal to the size difference.
+        For dimensions that shrink, chunks are truncated or removed to fit the new shape.
+
+        Parameters
+        ----------
+        new_shape : tuple[int, ...]
+            The new shape of the array. Must have the same number of dimensions as the chunk grid.
+
+        Returns
+        -------
+        Self
+            A new RectilinearChunkGrid instance with updated chunk shapes
+
+        Raises
+        ------
+        ValueError
+            If the number of dimensions in new_shape doesn't match the number of dimensions
+            in the chunk grid
+
+        Examples
+        --------
+        >>> grid = RectilinearChunkGrid(chunk_shapes=[[10, 20], [15, 15]])
+        >>> grid.update_shape((50, 40))  # Grow both dimensions
+        RectilinearChunkGrid(chunk_shapes=((10, 20, 20), (15, 15, 10)))
+
+        >>> grid = RectilinearChunkGrid(chunk_shapes=[[10, 20, 30], [25, 25]])
+        >>> grid.update_shape((25, 30))  # Shrink first dimension
+        RectilinearChunkGrid(chunk_shapes=((10, 20), (25, 25)))
+
+        Notes
+        -----
+        This method is automatically called when an array is resized. The chunk size
+        strategy for growing dimensions adds a single new chunk with size equal to
+        the growth amount. This may not be optimal for all use cases, and users may
+        want to manually adjust chunk shapes after resizing.
+        """
 
         if len(new_shape) != len(self.chunk_shapes):
             raise ValueError(

src/zarr/core/metadata/v3.py

@@ -104,8 +104,8 @@ def validate_codecs(codecs: tuple[Codec, ...], dtype: ZDType[TBaseDType, TBaseSc
     # we need to have special codecs if we are decoding vlen strings or bytestrings
     # TODO: use codec ID instead of class name
     codec_class_name = abc.__class__.__name__
-    # TODO: Fix typing here
-    if isinstance(dtype, VariableLengthUTF8) and not codec_class_name == "VLenUTF8Codec":
+    # TODO: Fix typing here - mypy cannot express the relationship between ZDType and VariableLengthUTF8
+    if isinstance(dtype, VariableLengthUTF8) and not codec_class_name == "VLenUTF8Codec":  # type: ignore[unreachable]
         raise ValueError(
             f"For string dtype, ArrayBytesCodec must be `VLenUTF8Codec`, got `{codec_class_name}`."
         )

src/zarr/testing/strategies.py

@@ -225,12 +225,17 @@ def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[
 
     For each dimension, generate a list of chunk sizes that sum to the dimension size.
     Sometimes uses uniform chunks, sometimes uses variable-sized chunks.
+
+    Note: We limit the number of chunks to max 20 per dimension to avoid performance issues
+    in property tests. With higher dimensions, the total chunk count grows multiplicatively.
     """
     chunk_shapes: list[list[int]] = []
     for size in shape:
         assert size > 0
         if size > 1:
-            nchunks = draw(st.integers(min_value=1, max_value=size - 1))
+            # Limit max chunks to 20 to avoid performance issues with large chunk grids
+            max_chunks = min(size - 1, 20)
+            nchunks = draw(st.integers(min_value=1, max_value=max_chunks))
             dividers = sorted(
                 draw(
                     st.lists(
@@ -486,10 +491,8 @@ def basic_indices(
         allow_ellipsis=allow_ellipsis,
     ).filter(
         lambda idxr: (
-            not (
-                is_negative_slice(idxr)
-                or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr))
-            )
+            not is_negative_slice(idxr)
+            and not (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr))  # type: ignore[redundant-expr]
         )
     )
     if math.prod(shape) >= 3:
@@ -575,27 +578,26 @@ def chunk_paths(draw: st.DrawFn, ndim: int, numblocks: tuple[int, ...], subset:
 @st.composite
 def complex_chunk_grids(draw: st.DrawFn) -> RectilinearChunkGrid:
     ndim = draw(st.integers(min_value=1, max_value=3))
-    nchunks = draw(st.integers(min_value=10, max_value=100))
-    # Don't require unique chunk sizes - rectilinear grids can have repeated sizes
-    dim_chunks = st.lists(
-        st.integers(min_value=1, max_value=10), min_size=nchunks, max_size=nchunks
-    )
+    # Limit to 5-10 chunks per dimension with small sizes to keep array size reasonable
+    # Max array size: 10 chunks * 5 size = 50 elements per dim, 50^3 = 125k elements max
+    nchunks = draw(st.integers(min_value=5, max_value=10))
+    # Keep chunk sizes small (1-5) to avoid creating huge arrays
+    dim_chunks = st.lists(st.integers(min_value=1, max_value=5), min_size=nchunks, max_size=nchunks)
     if draw(st.booleans()):
         event("using RectilinearChunkGrid")
         chunk_shapes = draw(st.lists(dim_chunks, min_size=ndim, max_size=ndim))
         return RectilinearChunkGrid(chunk_shapes=chunk_shapes)
 
     else:
         event("using RectilinearChunkGrid (run length encoded)")
-        # For RLE, we need to carefully control the total expanded chunks
-        # to avoid creating arrays that are too large
-        # Use a small number of RLE entries with small repeat counts
-        num_rle_entries = draw(st.integers(min_value=5, max_value=20))
+        # For RLE, keep total expanded chunks small: 3-5 entries * 2-3 repeats = 6-15 chunks
+        # With chunk sizes 1-5, max array size: 15 * 5 = 75 per dim, 75^3 = 421k elements max
+        num_rle_entries = draw(st.integers(min_value=3, max_value=5))
         chunk_shapes_rle = [
             [
                 [
-                    draw(st.integers(min_value=1, max_value=10)),  # chunk size
-                    draw(st.integers(min_value=1, max_value=3)),  # repeat count
+                    draw(st.integers(min_value=1, max_value=5)),  # chunk size
+                    draw(st.integers(min_value=2, max_value=3)),  # repeat count
                 ]
                 for _ in range(num_rle_entries)
             ]
@@ -613,7 +615,7 @@ def complex_chunked_arrays(
     draw: st.DrawFn,
     *,
     stores: st.SearchStrategy[StoreLike] = stores,
-) -> Array:
+) -> AnyArray:
     store = draw(stores, label="store")
     chunks = draw(complex_chunk_grids(), label="chunk grid")
     assert isinstance(chunks, RectilinearChunkGrid)

tests/test_array.py

@@ -67,6 +67,7 @@
 from zarr.core.group import AsyncGroup
 from zarr.core.indexing import BasicIndexer, _iter_grid, _iter_regions
 from zarr.core.metadata.v2 import ArrayV2Metadata
+from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.core.sync import sync
 from zarr.errors import (
     ContainsArrayError,