ChunkMetadata -> ArraySpec

Author: normanrz

zarr/v3/abc/codec.py

@@ -11,7 +11,7 @@
 
 if TYPE_CHECKING:
     from zarr.v3.metadata import (
-        ChunkMetadata,
+        ArraySpec,
         ArrayMetadata,
         DataType,
         CodecMetadata,
@@ -33,10 +33,10 @@ def from_metadata(cls, codec_metadata: CodecMetadata) -> Codec:
         pass
 
     @abstractmethod
-    def compute_encoded_size(self, input_byte_length: int, chunk_metadata: ChunkMetadata) -> int:
+    def compute_encoded_size(self, input_byte_length: int, chunk_metadata: ArraySpec) -> int:
         pass
 
-    def resolve_metadata(self, chunk_metadata: ChunkMetadata) -> ChunkMetadata:
+    def resolve_metadata(self, chunk_metadata: ArraySpec) -> ArraySpec:
         return chunk_metadata
 
     def evolve(self, *, ndim: int, data_type: DataType) -> Codec:
@@ -51,7 +51,7 @@ class ArrayArrayCodec(Codec):
     async def decode(
         self,
         chunk_array: np.ndarray,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
         pass
@@ -60,7 +60,7 @@ async def decode(
     async def encode(
         self,
         chunk_array: np.ndarray,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[np.ndarray]:
         pass
@@ -71,7 +71,7 @@ class ArrayBytesCodec(Codec):
     async def decode(
         self,
         chunk_array: BytesLike,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
         pass
@@ -80,7 +80,7 @@ async def decode(
     async def encode(
         self,
         chunk_array: np.ndarray,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         pass
@@ -92,7 +92,7 @@ async def decode_partial(
         self,
         store_path: StorePath,
         selection: SliceSelection,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[np.ndarray]:
         pass
@@ -105,7 +105,7 @@ async def encode_partial(
         store_path: StorePath,
         chunk_array: np.ndarray,
         selection: SliceSelection,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> None:
         pass
@@ -116,7 +116,7 @@ class BytesBytesCodec(Codec):
     async def decode(
         self,
         chunk_array: BytesLike,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> BytesLike:
         pass
@@ -125,7 +125,7 @@ async def decode(
     async def encode(
         self,
         chunk_array: BytesLike,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         pass

zarr/v3/array.py

@@ -33,7 +33,7 @@
 from zarr.v3.indexing import BasicIndexer, all_chunk_coords, is_total_slice
 from zarr.v3.metadata import (
     ArrayMetadata,
-    ChunkMetadata,
+    ArraySpec,
     DataType,
     DefaultChunkKeyEncodingConfigurationMetadata,
     DefaultChunkKeyEncodingMetadata,
@@ -381,7 +381,7 @@ async def _write_chunk(
             await self._write_chunk_to_store(store_path, chunk_array, chunk_metadata)
 
     async def _write_chunk_to_store(
-        self, store_path: StorePath, chunk_array: np.ndarray, chunk_metadata: ChunkMetadata
+        self, store_path: StorePath, chunk_array: np.ndarray, chunk_metadata: ArraySpec
     ):
         if np.all(chunk_array == self.metadata.fill_value):
             # chunks that only contain fill_value will be removed

zarr/v3/codecs/__init__.py

@@ -20,7 +20,7 @@
 from zarr.v3.metadata import CodecMetadata, ShardingCodecIndexLocation, RuntimeConfiguration
 
 if TYPE_CHECKING:
-    from zarr.v3.metadata import ArrayMetadata, ChunkMetadata
+    from zarr.v3.metadata import ArrayMetadata, ArraySpec
     from zarr.v3.codecs.sharding import ShardingCodecMetadata
     from zarr.v3.codecs.blosc import BloscCodecMetadata
     from zarr.v3.codecs.bytes import BytesCodecMetadata
@@ -31,8 +31,8 @@
 
 
 def _find_array_bytes_codec(
-    codecs: Iterable[Tuple[Codec, ChunkMetadata]]
-) -> Tuple[ArrayBytesCodec, ChunkMetadata]:
+    codecs: Iterable[Tuple[Codec, ArraySpec]]
+) -> Tuple[ArrayBytesCodec, ArraySpec]:
     for codec, chunk_metadata in codecs:
         if isinstance(codec, ArrayBytesCodec):
             return (codec, chunk_metadata)
@@ -88,16 +88,16 @@ def validate(self, array_metadata: ArrayMetadata) -> None:
             )
 
     def _codecs_with_resolved_metadata(
-        self, chunk_metadata: ChunkMetadata
-    ) -> Iterator[Tuple[Codec, ChunkMetadata]]:
+        self, chunk_metadata: ArraySpec
+    ) -> Iterator[Tuple[Codec, ArraySpec]]:
         for codec in self.codecs:
             yield (codec, chunk_metadata)
             chunk_metadata = codec.resolve_metadata(chunk_metadata)
 
     async def decode(
         self,
         chunk_bytes: BytesLike,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
         codecs = list(self._codecs_with_resolved_metadata(chunk_metadata))[::-1]
@@ -122,7 +122,7 @@ async def decode(
     async def encode(
         self,
         chunk_array: np.ndarray,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         codecs = list(self._codecs_with_resolved_metadata(chunk_metadata))
@@ -155,7 +155,7 @@ async def encode(
 
         return chunk_bytes
 
-    def compute_encoded_size(self, byte_length: int, chunk_metadata: ChunkMetadata) -> int:
+    def compute_encoded_size(self, byte_length: int, chunk_metadata: ArraySpec) -> int:
         for codec in self.codecs:
             byte_length = codec.compute_encoded_size(byte_length, chunk_metadata)
             chunk_metadata = codec.resolve_metadata(chunk_metadata)

zarr/v3/codecs/blosc.py

@@ -19,7 +19,7 @@
 from zarr.v3.common import BytesLike, to_thread
 
 if TYPE_CHECKING:
-    from zarr.v3.metadata import ChunkMetadata, CodecMetadata, DataType, RuntimeConfiguration
+    from zarr.v3.metadata import ArraySpec, CodecMetadata, DataType, RuntimeConfiguration
 
 
 BloscShuffle = Literal["noshuffle", "shuffle", "bitshuffle"]
@@ -86,21 +86,21 @@ def get_blosc_codec(self) -> Blosc:
     async def decode(
         self,
         chunk_bytes: bytes,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> BytesLike:
         return await to_thread(self.get_blosc_codec().decode, chunk_bytes)
 
     async def encode(
         self,
         chunk_bytes: bytes,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         chunk_array = np.frombuffer(chunk_bytes, dtype=chunk_metadata.dtype)
         return await to_thread(self.get_blosc_codec().encode, chunk_array)
 
-    def compute_encoded_size(self, _input_byte_length: int, _chunk_metadata: ChunkMetadata) -> int:
+    def compute_encoded_size(self, _input_byte_length: int, _chunk_metadata: ArraySpec) -> int:
         raise NotImplementedError
 
 

zarr/v3/codecs/bytes.py

@@ -15,7 +15,7 @@
 from zarr.v3.common import BytesLike
 
 if TYPE_CHECKING:
-    from zarr.v3.metadata import CodecMetadata, ChunkMetadata, ArrayMetadata, RuntimeConfiguration
+    from zarr.v3.metadata import CodecMetadata, ArraySpec, ArrayMetadata, RuntimeConfiguration
 
 
 Endian = Literal["big", "little"]
@@ -64,7 +64,7 @@ def _get_byteorder(self, array: np.ndarray) -> Endian:
     async def decode(
         self,
         chunk_bytes: BytesLike,
-        chunk_metadata: ChunkMetadata,
+        chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> np.ndarray:
         if chunk_metadata.dtype.itemsize > 0:
@@ -87,7 +87,7 @@ async def decode(
     async def encode(
         self,
         chunk_array: np.ndarray,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         if chunk_array.dtype.itemsize > 1:
@@ -97,7 +97,7 @@ async def encode(
                 chunk_array = chunk_array.astype(new_dtype)
         return chunk_array.tobytes()
 
-    def compute_encoded_size(self, input_byte_length: int, _chunk_metadata: ChunkMetadata) -> int:
+    def compute_encoded_size(self, input_byte_length: int, _chunk_metadata: ArraySpec) -> int:
         return input_byte_length
 
 

zarr/v3/codecs/crc32c_.py

@@ -16,7 +16,7 @@
 from zarr.v3.common import BytesLike
 
 if TYPE_CHECKING:
-    from zarr.v3.metadata import ChunkMetadata, CodecMetadata, RuntimeConfiguration
+    from zarr.v3.metadata import ArraySpec, CodecMetadata, RuntimeConfiguration
 
 
 @frozen
@@ -40,7 +40,7 @@ def get_metadata_class(cls) -> Type[Crc32cCodecMetadata]:
     async def decode(
         self,
         chunk_bytes: bytes,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> BytesLike:
         crc32_bytes = chunk_bytes[-4:]
@@ -52,12 +52,12 @@ async def decode(
     async def encode(
         self,
         chunk_bytes: bytes,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         return chunk_bytes + np.uint32(crc32c(chunk_bytes)).tobytes()
 
-    def compute_encoded_size(self, input_byte_length: int, _chunk_metadata: ChunkMetadata) -> int:
+    def compute_encoded_size(self, input_byte_length: int, _chunk_metadata: ArraySpec) -> int:
         return input_byte_length + 4
 
 

zarr/v3/codecs/gzip.py

@@ -15,7 +15,7 @@
 from zarr.v3.common import BytesLike, to_thread
 
 if TYPE_CHECKING:
-    from zarr.v3.metadata import ChunkMetadata, CodecMetadata, RuntimeConfiguration
+    from zarr.v3.metadata import ArraySpec, CodecMetadata, RuntimeConfiguration
 
 
 @frozen
@@ -47,23 +47,23 @@ def get_metadata_class(cls) -> Type[GzipCodecMetadata]:
     async def decode(
         self,
         chunk_bytes: bytes,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> BytesLike:
         return await to_thread(GZip(self.configuration.level).decode, chunk_bytes)
 
     async def encode(
         self,
         chunk_bytes: bytes,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
         _runtime_configuration: RuntimeConfiguration,
     ) -> Optional[BytesLike]:
         return await to_thread(GZip(self.configuration.level).encode, chunk_bytes)
 
     def compute_encoded_size(
         self,
         _input_byte_length: int,
-        _chunk_metadata: ChunkMetadata,
+        _chunk_metadata: ArraySpec,
     ) -> int:
         raise NotImplementedError