revert removal of metadata chunk grid attribute

Author: d-v-b

src/zarr/core/array.py

@@ -5,7 +5,6 @@
 from asyncio import gather
 from collections.abc import Iterable
 from dataclasses import dataclass, field, replace
-from functools import cached_property
 from itertools import starmap
 from logging import getLogger
 from typing import (
@@ -32,7 +31,7 @@
 from zarr.codecs._v2 import V2Codec
 from zarr.codecs.bytes import BytesCodec
 from zarr.core._info import ArrayInfo
-from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArraySpec, parse_array_config
+from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, parse_array_config
 from zarr.core.attributes import Attributes
 from zarr.core.buffer import (
     BufferPrototype,
@@ -42,7 +41,7 @@
     default_buffer_prototype,
 )
 from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype
-from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid, _auto_partition, normalize_chunks
+from zarr.core.chunk_grids import RegularChunkGrid, _auto_partition, normalize_chunks
 from zarr.core.chunk_key_encodings import (
     ChunkKeyEncoding,
     ChunkKeyEncodingLike,
@@ -951,13 +950,6 @@ def chunks(self) -> ChunkCoords:
         """
         return self.metadata.chunks
 
-    @cached_property
-    def chunk_grid(self) -> ChunkGrid:
-        if self.metadata.zarr_format == 2:
-            return RegularChunkGrid(chunk_shape=self.chunks)
-        else:
-            return self.metadata.chunk_grid
-
     @property
     def shards(self) -> ChunkCoords | None:
         """Returns the shard shape of the Array.
@@ -1281,20 +1273,6 @@ def nbytes(self) -> int:
         """
         return self.size * self.dtype.itemsize
 
-    def get_chunk_spec(
-        self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype
-    ) -> ArraySpec:
-        assert isinstance(self.chunk_grid, RegularChunkGrid), (
-            "Currently, only regular chunk grid is supported"
-        )
-        return ArraySpec(
-            shape=self.chunk_grid.chunk_shape,
-            dtype=self._zdtype,
-            fill_value=self.metadata.fill_value,
-            config=array_config,
-            prototype=prototype,
-        )
-
     async def _get_selection(
         self,
         indexer: Indexer,
@@ -1334,7 +1312,7 @@ async def _get_selection(
                 [
                     (
                         self.store_path / self.metadata.encode_chunk_key(chunk_coords),
-                        self.get_chunk_spec(chunk_coords, _config, prototype=prototype),
+                        self.metadata.get_chunk_spec(chunk_coords, _config, prototype=prototype),
                         chunk_selection,
                         out_selection,
                         is_complete_chunk,
@@ -1389,7 +1367,7 @@ async def getitem(
         indexer = BasicIndexer(
             selection,
             shape=self.metadata.shape,
-            chunk_grid=self.chunk_grid,
+            chunk_grid=self.metadata.chunk_grid,
         )
         return await self._get_selection(indexer, prototype=prototype)
 
@@ -1464,7 +1442,7 @@ async def _set_selection(
             [
                 (
                     self.store_path / self.metadata.encode_chunk_key(chunk_coords),
-                    self.get_chunk_spec(chunk_coords, _config, prototype),
+                    self.metadata.get_chunk_spec(chunk_coords, _config, prototype),
                     chunk_selection,
                     out_selection,
                     is_complete_chunk,
@@ -1519,7 +1497,7 @@ async def setitem(
         indexer = BasicIndexer(
             selection,
             shape=self.metadata.shape,
-            chunk_grid=self.chunk_grid,
+            chunk_grid=self.metadata.chunk_grid,
         )
         return await self._set_selection(indexer, value, prototype=prototype)
 
@@ -1556,8 +1534,8 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True)
 
         if delete_outside_chunks:
             # Remove all chunks outside of the new shape
-            old_chunk_coords = set(self.chunk_grid.all_chunk_coords(self.metadata.shape))
-            new_chunk_coords = set(self.chunk_grid.all_chunk_coords(new_shape))
+            old_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(self.metadata.shape))
+            new_chunk_coords = set(self.metadata.chunk_grid.all_chunk_coords(new_shape))
 
             async def _delete_key(key: str) -> None:
                 await (self.store_path / key).delete()
@@ -2687,7 +2665,7 @@ def get_basic_selection(
             prototype = default_buffer_prototype()
         return sync(
             self._async_array._get_selection(
-                BasicIndexer(selection, self.shape, self._async_array.chunk_grid),
+                BasicIndexer(selection, self.shape, self.metadata.chunk_grid),
                 out=out,
                 fields=fields,
                 prototype=prototype,
@@ -2787,7 +2765,7 @@ def set_basic_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = BasicIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = BasicIndexer(selection, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 
     @_deprecate_positional_args
@@ -2908,7 +2886,7 @@ def get_orthogonal_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = OrthogonalIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
                 indexer=indexer, out=out, fields=fields, prototype=prototype
@@ -3021,7 +2999,7 @@ def set_orthogonal_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = OrthogonalIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = OrthogonalIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype)
         )
@@ -3102,7 +3080,7 @@ def get_mask_selection(
 
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = MaskIndexer(mask, self.shape, self._async_array.chunk_grid)
+        indexer = MaskIndexer(mask, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
                 indexer=indexer, out=out, fields=fields, prototype=prototype
@@ -3185,7 +3163,7 @@ def set_mask_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = MaskIndexer(mask, self.shape, self._async_array.chunk_grid)
+        indexer = MaskIndexer(mask, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 
     @_deprecate_positional_args
@@ -3266,7 +3244,7 @@ def get_coordinate_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = CoordinateIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = CoordinateIndexer(selection, self.shape, self.metadata.chunk_grid)
         out_array = sync(
             self._async_array._get_selection(
                 indexer=indexer, out=out, fields=fields, prototype=prototype
@@ -3352,7 +3330,7 @@ def set_coordinate_selection(
         if prototype is None:
             prototype = default_buffer_prototype()
         # setup indexer
-        indexer = CoordinateIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = CoordinateIndexer(selection, self.shape, self.metadata.chunk_grid)
 
         # handle value - need ndarray-like flatten value
         if not is_scalar(value, self.dtype):
@@ -3468,7 +3446,7 @@ def get_block_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = BlockIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = BlockIndexer(selection, self.shape, self.metadata.chunk_grid)
         return sync(
             self._async_array._get_selection(
                 indexer=indexer, out=out, fields=fields, prototype=prototype
@@ -3562,7 +3540,7 @@ def set_block_selection(
         """
         if prototype is None:
             prototype = default_buffer_prototype()
-        indexer = BlockIndexer(selection, self.shape, self._async_array.chunk_grid)
+        indexer = BlockIndexer(selection, self.shape, self.metadata.chunk_grid)
         sync(self._async_array._set_selection(indexer, value, fields=fields, prototype=prototype))
 
     @property

src/zarr/core/metadata/v2.py

@@ -3,11 +3,13 @@
 import base64
 import warnings
 from collections.abc import Iterable, Sequence
+from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypeAlias, TypedDict
 
 import numcodecs.abc
 
 from zarr.abc.metadata import Metadata
+from zarr.core.chunk_grids import RegularChunkGrid
 from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, TDType_co, TScalar_co, ZDType
 
@@ -103,6 +105,10 @@ def __init__(
     def ndim(self) -> int:
         return len(self.shape)
 
+    @cached_property
+    def chunk_grid(self) -> RegularChunkGrid:
+        return RegularChunkGrid(chunk_shape=self.chunks)
+
     @property
     def shards(self) -> ChunkCoords | None:
         return None

src/zarr/core/metadata/v3.py

@@ -269,6 +269,20 @@ def inner_codecs(self) -> tuple[Codec, ...]:
                 return self.codecs[0].codecs
         return self.codecs
 
+    def get_chunk_spec(
+        self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype
+    ) -> ArraySpec:
+        assert isinstance(self.chunk_grid, RegularChunkGrid), (
+            "Currently, only regular chunk grid is supported"
+        )
+        return ArraySpec(
+            shape=self.chunk_grid.chunk_shape,
+            dtype=self.dtype,
+            fill_value=self.fill_value,
+            config=array_config,
+            prototype=prototype,
+        )
+
     def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str:
         return self.chunk_key_encoding.encode_chunk_key(chunk_coords)
 

tests/test_array.py

@@ -1363,7 +1363,7 @@ async def test_with_data(impl: Literal["sync", "async"], store: Store) -> None:
         elif impl == "async":
             arr = await create_array(store, name=name, data=data, zarr_format=3)
             stored = await arr._get_selection(
-                BasicIndexer(..., shape=arr.shape, chunk_grid=arr.chunk_grid),
+                BasicIndexer(..., shape=arr.shape, chunk_grid=arr.metadata.chunk_grid),
                 prototype=default_buffer_prototype(),
             )
         else:

tests/test_group.py

@@ -1007,7 +1007,7 @@ async def test_asyncgroup_create_array(
     assert subnode.dtype == dtype
     # todo: fix the type annotation of array.metadata.chunk_grid so that we get some autocomplete
     # here.
-    assert subnode.chunk_grid.chunk_shape == chunk_shape
+    assert subnode.metadata.chunk_grid.chunk_shape == chunk_shape
     assert subnode.metadata.zarr_format == zarr_format