adapt Array.info

Author: normanrz

src/zarr/core/_info.py

@@ -5,7 +5,7 @@
 import numcodecs.abc
 import numpy as np
 
-from zarr.abc.codec import Codec
+from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
 from zarr.core.common import ZarrFormat
 from zarr.core.metadata.v3 import DataType
 
@@ -85,9 +85,9 @@ class ArrayInfo:
     _order: Literal["C", "F"]
     _read_only: bool
     _store_type: str
-    _compressor: numcodecs.abc.Codec | None = None
-    _filters: tuple[numcodecs.abc.Codec, ...] | None = None
-    _codecs: tuple[Codec, ...] | None = None
+    _filters: tuple[numcodecs.abc.Codec, ...] | tuple[ArrayArrayCodec, ...] = ()
+    _serializer: ArrayBytesCodec | None = None
+    _compressors: tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec, ...] = ()
     _count_bytes: int | None = None
     _count_bytes_stored: int | None = None
     _count_chunks_initialized: int | None = None
@@ -113,14 +113,13 @@ def __repr__(self) -> str:
         if self._chunk_shape is None:
             # for non-regular chunk grids
             kwargs["chunk_shape"] = "<variable>"
-        if self._compressor is not None:
-            template += "\nCompressor         : {_compressor}"
 
-        if self._filters is not None:
+        if len(self._filters) > 0:
             template += "\nFilters            : {_filters}"
-
-        if self._codecs is not None:
-            template += "\nCodecs             : {_codecs}"
+        if self._serializer is not None:
+            template += "\nSerializer         : {_serializer}"
+        if len(self._compressors) > 0:
+            template += "\nCompressors        : {_compressors}"
 
         if self._count_bytes is not None:
             template += "\nNo. bytes          : {_count_bytes}"
@@ -139,5 +138,8 @@ def __repr__(self) -> str:
             kwargs["_storage_ratio"] = f"{self._count_bytes / self._count_bytes_stored:.1f}"
 
         if self._count_chunks_initialized is not None:
-            template += "\nChunks Initialized : {_count_chunks_initialized}"
+            if self._shard_shape is not None:
+                template += "\nShards Initialized : {_count_chunks_initialized}"
+            else:
+                template += "\nChunks Initialized : {_count_chunks_initialized}"
         return template.format(**kwargs)

src/zarr/core/array.py

@@ -924,7 +924,9 @@ def filters(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[ArrayArrayCodec, ..
                 return ()
             return filters
 
-        return tuple(codec for codec in self.metadata.codecs if isinstance(codec, ArrayArrayCodec))
+        return tuple(
+            codec for codec in self.metadata.inner_codecs if isinstance(codec, ArrayArrayCodec)
+        )
 
     @property
     def serializer(self) -> ArrayBytesCodec | None:
@@ -934,7 +936,9 @@ def serializer(self) -> ArrayBytesCodec | None:
         if self.metadata.zarr_format == 2:
             return None
 
-        return next(codec for codec in self.metadata.codecs if isinstance(codec, ArrayBytesCodec))
+        return next(
+            codec for codec in self.metadata.inner_codecs if isinstance(codec, ArrayBytesCodec)
+        )
 
     @property
     @deprecated("Use AsyncArray.compressors instead.")
@@ -961,7 +965,9 @@ def compressors(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec
                 return (self.metadata.compressor,)
             return ()
 
-        return tuple(codec for codec in self.metadata.codecs if isinstance(codec, BytesBytesCodec))
+        return tuple(
+            codec for codec in self.metadata.inner_codecs if isinstance(codec, BytesBytesCodec)
+        )
 
     @property
     def dtype(self) -> np.dtype[Any]:
@@ -1613,31 +1619,26 @@ async def info_complete(self) -> Any:
     def _info(
         self, count_chunks_initialized: int | None = None, count_bytes_stored: int | None = None
     ) -> Any:
-        kwargs: dict[str, Any] = {}
-        if self.metadata.zarr_format == 2:
-            assert isinstance(self.metadata, ArrayV2Metadata)
-            if self.metadata.compressor is not None:
-                kwargs["_compressor"] = self.metadata.compressor
-            if self.metadata.filters is not None:
-                kwargs["_filters"] = self.metadata.filters
-            kwargs["_data_type"] = self.metadata.dtype
-            kwargs["_chunk_shape"] = self.metadata.chunks
+        _data_type: np.dtype[Any] | DataType
+        if isinstance(self.metadata, ArrayV2Metadata):
+            _data_type = self.metadata.dtype
         else:
-            kwargs["_codecs"] = self.metadata.codecs
-            kwargs["_data_type"] = self.metadata.data_type
-            kwargs["_chunk_shape"] = self.chunks
-            kwargs["_shard_shape"] = self.shards
-
+            _data_type = self.metadata.data_type
         return ArrayInfo(
             _zarr_format=self.metadata.zarr_format,
+            _data_type=_data_type,
             _shape=self.shape,
             _order=self.order,
+            _shard_shape=self.shards,
+            _chunk_shape=self.chunks,
             _read_only=self.read_only,
+            _compressors=self.compressors,
+            _filters=self.filters,
+            _serializer=self.serializer,
             _store_type=type(self.store_path.store).__name__,
             _count_bytes=self.nbytes,
             _count_bytes_stored=count_bytes_stored,
             _count_chunks_initialized=count_chunks_initialized,
-            **kwargs,
         )
 
 

src/zarr/core/metadata/v3.py

@@ -81,17 +81,26 @@ def parse_codecs(data: object) -> tuple[Codec, ...]:
     return out
 
 
-def validate_codecs(codecs: tuple[Codec, ...], dtype: DataType) -> None:
-    """Check that the codecs are valid for the given dtype"""
-
+def validate_array_bytes_codec(codecs: tuple[Codec, ...]) -> ArrayBytesCodec:
     # ensure that we have at least one ArrayBytesCodec
     abcs: list[ArrayBytesCodec] = [codec for codec in codecs if isinstance(codec, ArrayBytesCodec)]
     if len(abcs) == 0:
         raise ValueError("At least one ArrayBytesCodec is required.")
     elif len(abcs) > 1:
         raise ValueError("Only one ArrayBytesCodec is allowed.")
 
-    abc = abcs[0]
+    return abcs[0]
+
+
+def validate_codecs(codecs: tuple[Codec, ...], dtype: DataType) -> None:
+    """Check that the codecs are valid for the given dtype"""
+    from zarr.codecs.sharding import ShardingCodec
+
+    abc = validate_array_bytes_codec(codecs)
+
+    # Recursively resolve array-bytes codecs within sharding codecs
+    while isinstance(abc, ShardingCodec):
+        abc = validate_array_bytes_codec(abc.codecs)
 
     # we need to have special codecs if we are decoding vlen strings or bytestrings
     # TODO: use codec ID instead of class name
@@ -330,6 +339,15 @@ def shards(self) -> ChunkCoords | None:
         )
         raise NotImplementedError(msg)
 
+    @property
+    def inner_codecs(self) -> tuple[Codec, ...]:
+        if isinstance(self.chunk_grid, RegularChunkGrid):
+            from zarr.codecs.sharding import ShardingCodec
+
+            if len(self.codecs) == 1 and isinstance(self.codecs[0], ShardingCodec):
+                return self.codecs[0].codecs
+        return self.codecs
+
     def get_chunk_spec(
         self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype
     ) -> ArraySpec:

tests/test_array.py

@@ -20,7 +20,6 @@
     VLenUTF8Codec,
     ZstdCodec,
 )
-from zarr.codecs.sharding import ShardingCodec
 from zarr.core._info import ArrayInfo
 from zarr.core.array import (
     CompressorsLike,
@@ -494,7 +493,7 @@ def test_info_v2(self, chunks: tuple[int, int], shards: tuple[int, int] | None)
             _read_only=False,
             _store_type="MemoryStore",
             _count_bytes=512,
-            _compressor=numcodecs.Zstd(),
+            _compressors=(numcodecs.Zstd(),),
         )
         assert result == expected
 
@@ -510,9 +509,8 @@ def test_info_v3(self, chunks: tuple[int, int], shards: tuple[int, int] | None)
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=(BytesCodec(), ZstdCodec())
-            if shards is None
-            else (ShardingCodec(chunk_shape=chunks, codecs=[BytesCodec(), ZstdCodec()]),),
+            _compressors=(ZstdCodec(),),
+            _serializer=BytesCodec(),
             _count_bytes=512,
         )
         assert result == expected
@@ -536,7 +534,7 @@ def test_info_complete(self, chunks: tuple[int, int], shards: tuple[int, int] |
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=(BytesCodec(),) if shards is None else (ShardingCodec(chunk_shape=chunks),),
+            _serializer=BytesCodec(),
             _count_bytes=512,
             _count_chunks_initialized=0,
             _count_bytes_stored=373 if shards is None else 578,  # the metadata?
@@ -572,7 +570,7 @@ async def test_info_v2_async(
             _read_only=False,
             _store_type="MemoryStore",
             _count_bytes=512,
-            _compressor=numcodecs.Zstd(),
+            _compressors=(numcodecs.Zstd(),),
         )
         assert result == expected
 
@@ -596,9 +594,8 @@ async def test_info_v3_async(
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=(BytesCodec(), ZstdCodec())
-            if shards is None
-            else (ShardingCodec(chunk_shape=chunks, codecs=[BytesCodec(), ZstdCodec()]),),
+            _compressors=(ZstdCodec(),),
+            _serializer=BytesCodec(),
             _count_bytes=512,
         )
         assert result == expected
@@ -624,7 +621,7 @@ async def test_info_complete_async(
             _order="C",
             _read_only=False,
             _store_type="MemoryStore",
-            _codecs=(BytesCodec(),) if shards is None else (ShardingCodec(chunk_shape=chunks),),
+            _serializer=BytesCodec(),
             _count_bytes=512,
             _count_chunks_initialized=0,
             _count_bytes_stored=373 if shards is None else 578,  # the metadata?
@@ -1125,16 +1122,24 @@ async def test_create_array_no_filters_compressors(
         ({"name": "transpose", "configuration": {"order": [0]}},),
     ],
 )
+@pytest.mark.parametrize(("chunks", "shards"), [((6,), None), ((3,), (6,))])
 async def test_create_array_v3_chunk_encoding(
-    store: MemoryStore, compressors: CompressorsLike, filters: FiltersLike, dtype: str
+    store: MemoryStore,
+    compressors: CompressorsLike,
+    filters: FiltersLike,
+    dtype: str,
+    chunks: tuple[int, ...],
+    shards: tuple[int, ...] | None,
 ) -> None:
     """
     Test various possibilities for the compressors and filters parameter to create_array
     """
     arr = await create_array(
         store=store,
         dtype=dtype,
-        shape=(10,),
+        shape=(12,),
+        chunks=chunks,
+        shards=shards,
         zarr_format=3,
         filters=filters,
         compressors=compressors,

tests/test_config.py

@@ -305,12 +305,12 @@ class NewCodec2(BytesCodec):
 @pytest.mark.parametrize(
     ("dtype", "expected_codecs"),
     [
-        ("int", [BytesCodec(), GzipCodec()]),
-        ("bytes", [VLenBytesCodec(), GzipCodec()]),
-        ("str", [VLenUTF8Codec(), GzipCodec()]),
+        ("int", (BytesCodec(), GzipCodec())),
+        ("bytes", (VLenBytesCodec(), GzipCodec())),
+        ("str", (VLenUTF8Codec(), GzipCodec())),
     ],
 )
-async def test_default_codecs(dtype: str, expected_codecs: list[Codec]) -> None:
+async def test_default_codecs(dtype: str, expected_codecs: tuple[Codec, ...]) -> None:
     with config.set(
         {
             "array.v3_default_codecs": {  # test setting non-standard codecs

tests/test_info.py

@@ -59,7 +59,7 @@ def test_array_info(zarr_format: ZarrFormat) -> None:
         _order="C",
         _read_only=True,
         _store_type="MemoryStore",
-        _codecs=(BytesCodec(),),
+        _serializer=BytesCodec(),
     )
     result = repr(info)
     assert result == textwrap.dedent(f"""\
@@ -71,7 +71,7 @@ def test_array_info(zarr_format: ZarrFormat) -> None:
         Order              : C
         Read-only          : True
         Store type         : MemoryStore
-        Codecs             : ({{'endian': <Endian.little: 'little'>}},)""")
+        Serializer         : {{'endian': <Endian.little: 'little'>}}""")
 
 
 @pytest.mark.parametrize("zarr_format", ZARR_FORMATS)
@@ -95,7 +95,7 @@ def test_array_info_complete(
         _order="C",
         _read_only=True,
         _store_type="MemoryStore",
-        _codecs=(BytesCodec(),),
+        _serializer=BytesCodec(),
         _count_bytes=count_bytes,
         _count_bytes_stored=count_bytes_stored,
         _count_chunks_initialized=count_chunks_initialized,
@@ -110,7 +110,7 @@ def test_array_info_complete(
         Order              : C
         Read-only          : True
         Store type         : MemoryStore
-        Codecs             : ({{'endian': <Endian.little: 'little'>}},)
+        Serializer         : {{'endian': <Endian.little: 'little'>}}
         No. bytes          : {count_bytes} ({count_bytes_formatted})
         No. bytes stored   : {count_bytes_stored_formatted}
         Storage ratio      : {storage_ratio_formatted}