Store array_shape on ChunkGrid instead of passing it per-method

Move array_shape from a per-method parameter on ChunkGrid to a stored
attribute, validated once at construction time. This simplifies call
sites (especially in indexing.py) by eliminating redundant array_shape
arguments throughout the codebase.

Key changes:
- Add abstract `array_shape` property to ChunkGrid ABC
- Add `_array_shape` field and `array_shape` property to RegularChunkGrid
- Remove `array_shape` parameter from all ChunkGrid abstract methods
- Update `from_dict` to accept `array_shape` keyword argument
- Update `update_shape` to return a new instance with updated array_shape
- Refactor indexing.py to use chunk_grid methods directly instead of
  raw arithmetic on chunk_shape
- Update all construction sites to pass array_shape

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: d-v-b

src/zarr/codecs/sharding.py

@@ -415,7 +415,7 @@ async def _decode_single(
         indexer = BasicIndexer(
             tuple(slice(0, s) for s in shard_shape),
             shape=shard_shape,
-            chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
+            chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape, array_shape=shard_shape),
         )
 
         # setup output array
@@ -461,7 +461,7 @@ async def _decode_partial_single(
         indexer = get_indexer(
             selection,
             shape=shard_shape,
-            chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
+            chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape, array_shape=shard_shape),
         )
 
         # setup output array
@@ -536,7 +536,7 @@ async def _encode_single(
             BasicIndexer(
                 tuple(slice(0, s) for s in shard_shape),
                 shape=shard_shape,
-                chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape),
+                chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape, array_shape=shard_shape),
             )
         )
 
@@ -585,7 +585,9 @@ async def _encode_partial_single(
 
         indexer = list(
             get_indexer(
-                selection, shape=shard_shape, chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape)
+                selection,
+                shape=shard_shape,
+                chunk_grid=RegularChunkGrid(chunk_shape=chunk_shape, array_shape=shard_shape),
             )
         )
 

src/zarr/core/array.py

@@ -774,7 +774,7 @@ def _create_metadata_v3(
         else:
             fill_value_parsed = fill_value
 
-        chunk_grid_parsed = RegularChunkGrid(chunk_shape=chunk_shape)
+        chunk_grid_parsed = RegularChunkGrid(chunk_shape=chunk_shape, array_shape=shape)
         return ArrayV3Metadata(
             shape=shape,
             data_type=dtype,
@@ -4694,7 +4694,9 @@ async def init_array(
             sharding_codec.validate(
                 shape=chunk_shape_parsed,
                 dtype=zdtype,
-                chunk_grid=RegularChunkGrid(chunk_shape=shard_shape_parsed),
+                chunk_grid=RegularChunkGrid(
+                    chunk_shape=shard_shape_parsed, array_shape=chunk_shape_parsed
+                ),
             )
             codecs_out = (sharding_codec,)
             chunks_out = shard_shape_parsed
@@ -5995,8 +5997,8 @@ async def _resize(
 
     if delete_outside_chunks and not only_growing:
         # Remove all chunks outside of the new shape
-        old_chunk_coords = set(array.metadata.chunk_grid.all_chunk_coords(array.metadata.shape))
-        new_chunk_coords = set(array.metadata.chunk_grid.all_chunk_coords(new_shape))
+        old_chunk_coords = set(array.metadata.chunk_grid.all_chunk_coords())
+        new_chunk_coords = set(new_metadata.chunk_grid.all_chunk_coords())
 
         async def _delete_key(key: str) -> None:
             await (array.store_path / key).delete()

src/zarr/core/chunk_grids/__init__.py

@@ -18,6 +18,8 @@
 
 def parse_chunk_grid(
     data: dict[str, JSON] | ChunkGrid | NamedConfig[str, Any],
+    *,
+    array_shape: tuple[int, ...],
 ) -> ChunkGrid:
     """Parse a chunk grid from a dictionary, returning existing ChunkGrid instances as-is.
 
@@ -28,6 +30,8 @@ def parse_chunk_grid(
     data : dict[str, JSON] | ChunkGrid | NamedConfig[str, Any]
         Either a ChunkGrid instance (returned as-is) or a dictionary with
         'name' and 'configuration' keys.
+    array_shape : tuple[int, ...]
+        The shape of the array this chunk grid is bound to.
 
     Returns
     -------
@@ -46,7 +50,7 @@ def parse_chunk_grid(
         chunk_grid_cls = get_chunk_grid_class(name_parsed)
     except KeyError as e:
         raise ValueError(f"Unknown chunk grid. Got {name_parsed}.") from e
-    return chunk_grid_cls.from_dict(data)  # type: ignore[arg-type]
+    return chunk_grid_cls.from_dict(data, array_shape=array_shape)  # type: ignore[arg-type, call-arg]
 
 
 __all__ = [

src/zarr/core/chunk_grids/common.py

@@ -3,9 +3,8 @@
 import math
 import numbers
 import warnings
-from abc import abstractmethod
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Literal
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Literal, Self
 
 import numpy as np
 import numpy.typing as npt
@@ -16,144 +15,63 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from typing import Self
 
     from zarr.core.array import ShardsLike
     from zarr.core.common import JSON
 
 
-@dataclass(frozen=True)
-class ChunkGrid(Metadata):
+class ChunkGrid(ABC, Metadata):
+    @property
+    @abstractmethod
+    def array_shape(self) -> tuple[int, ...]:
+        """The shape of the array this chunk grid is bound to."""
+        ...
+
     @abstractmethod
     def to_dict(self) -> dict[str, JSON]: ...
 
     @abstractmethod
     def update_shape(self, new_shape: tuple[int, ...]) -> Self:
-        pass
+        """Return a new ChunkGrid with the given array_shape."""
+        ...
 
     @abstractmethod
-    def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]:
-        pass
+    def all_chunk_coords(self) -> Iterator[tuple[int, ...]]: ...
 
     @abstractmethod
-    def get_nchunks(self, array_shape: tuple[int, ...]) -> int:
-        pass
+    def get_nchunks(self) -> int: ...
 
     @abstractmethod
-    def get_chunk_shape(
-        self, array_shape: tuple[int, ...], chunk_coord: tuple[int, ...]
-    ) -> tuple[int, ...]:
-        """
-        Get the shape of a specific chunk.
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-        chunk_coord : tuple[int, ...]
-            Coordinates of the chunk in the chunk grid.
-
-        Returns
-        -------
-        tuple[int, ...]
-            Shape of the chunk at the given coordinates.
-        """
+    def get_chunk_shape(self, chunk_coord: tuple[int, ...]) -> tuple[int, ...]:
+        """Get the shape of a specific chunk."""
+        ...
 
     @abstractmethod
-    def get_chunk_start(
-        self, array_shape: tuple[int, ...], chunk_coord: tuple[int, ...]
-    ) -> tuple[int, ...]:
-        """
-        Get the starting position of a chunk in the array.
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-        chunk_coord : tuple[int, ...]
-            Coordinates of the chunk in the chunk grid.
-
-        Returns
-        -------
-        tuple[int, ...]
-            Starting position (offset) of the chunk in the array.
-        """
+    def get_chunk_start(self, chunk_coord: tuple[int, ...]) -> tuple[int, ...]:
+        """Get the starting position of a chunk in the array."""
+        ...
 
     @abstractmethod
-    def array_index_to_chunk_coord(
-        self, array_shape: tuple[int, ...], array_index: tuple[int, ...]
-    ) -> tuple[int, ...]:
-        """
-        Map an array index to the chunk coordinates that contain it.
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-        array_index : tuple[int, ...]
-            Index in the array.
-
-        Returns
-        -------
-        tuple[int, ...]
-            Coordinates of the chunk containing the array index.
-        """
+    def array_index_to_chunk_coord(self, array_index: tuple[int, ...]) -> tuple[int, ...]:
+        """Map an array index to the chunk coordinates that contain it."""
+        ...
 
     @abstractmethod
     def array_indices_to_chunk_dim(
-        self, array_shape: tuple[int, ...], dim: int, indices: npt.NDArray[np.intp]
+        self, dim: int, indices: npt.NDArray[np.intp]
     ) -> npt.NDArray[np.intp]:
-        """
-        Map an array of indices along one dimension to chunk coordinates (vectorized).
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-        dim : int
-            Dimension index.
-        indices : np.ndarray
-            Array of indices along the given dimension.
-
-        Returns
-        -------
-        np.ndarray
-            Array of chunk coordinates, same shape as indices.
-        """
+        """Map an array of indices along one dimension to chunk coordinates (vectorized)."""
+        ...
 
     @abstractmethod
-    def chunks_per_dim(self, array_shape: tuple[int, ...], dim: int) -> int:
-        """
-        Get the number of chunks along a specific dimension.
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-        dim : int
-            Dimension index.
-
-        Returns
-        -------
-        int
-            Number of chunks along the dimension.
-        """
+    def chunks_per_dim(self, dim: int) -> int:
+        """Get the number of chunks along a specific dimension."""
+        ...
 
     @abstractmethod
-    def get_chunk_grid_shape(self, array_shape: tuple[int, ...]) -> tuple[int, ...]:
-        """
-        Get the shape of the chunk grid (number of chunks along each dimension).
-
-        Parameters
-        ----------
-        array_shape : tuple[int, ...]
-            Shape of the full array.
-
-        Returns
-        -------
-        tuple[int, ...]
-            Number of chunks along each dimension.
-        """
+    def get_chunk_grid_shape(self) -> tuple[int, ...]:
+        """Get the shape of the chunk grid (number of chunks along each dimension)."""
+        ...
 
 
 def _guess_chunks(