Merge branch 'chore/handle-numcodecs-codecs' of github.com:d-v-b/zarr-python into feat/v2-v3-codecs

Author: d-v-b

src/zarr/codecs/__init__.py

@@ -4,10 +4,34 @@
 from zarr.codecs.bytes import BytesCodec, Endian
 from zarr.codecs.crc32c_ import Crc32cCodec
 from zarr.codecs.gzip import GzipCodec
+from zarr.codecs.numcodecs import (
+    BZ2,
+    CRC32,
+    CRC32C,
+    LZ4,
+    LZMA,
+    ZFPY,
+    Adler32,
+    AsType,
+    BitRound,
+    Blosc,
+    Delta,
+    FixedScaleOffset,
+    Fletcher32,
+    GZip,
+    JenkinsLookup3,
+    PackBits,
+    PCodec,
+    Quantize,
+    Shuffle,
+    Zlib,
+    Zstd,
+)
 from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
 from zarr.codecs.transpose import TransposeCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
 from zarr.codecs.zstd import ZstdCodec
+from zarr.registry import register_codec
 
 __all__ = [
     "BloscCname",
@@ -24,3 +48,46 @@
     "VLenUTF8Codec",
     "ZstdCodec",
 ]
+
+register_codec("blosc", BloscCodec)
+register_codec("bytes", BytesCodec)
+
+# compatibility with earlier versions of ZEP1
+register_codec("endian", BytesCodec)
+register_codec("crc32c", Crc32cCodec)
+register_codec("gzip", GzipCodec)
+register_codec("sharding_indexed", ShardingCodec)
+register_codec("zstd", ZstdCodec)
+register_codec("vlen-utf8", VLenUTF8Codec)
+register_codec("vlen-bytes", VLenBytesCodec)
+register_codec("transpose", TransposeCodec)
+
+# Register all the codecs formerly contained in numcodecs.zarr3
+
+register_codec("numcodecs.bz2", BZ2, qualname="zarr.codecs.numcodecs.BZ2")
+register_codec("numcodecs.crc32", CRC32, qualname="zarr.codecs.numcodecs.CRC32")
+register_codec("numcodecs.crc32c", CRC32C, qualname="zarr.codecs.numcodecs.CRC32C")
+register_codec("numcodecs.lz4", LZ4, qualname="zarr.codecs.numcodecs.LZ4")
+register_codec("numcodecs.lzma", LZMA, qualname="zarr.codecs.numcodecs.LZMA")
+register_codec("numcodecs.zfpy", ZFPY, qualname="zarr.codecs.numcodecs.ZFPY")
+register_codec("numcodecs.adler32", Adler32, qualname="zarr.codecs.numcodecs.Adler32")
+register_codec("numcodecs.astype", AsType, qualname="zarr.codecs.numcodecs.AsType")
+register_codec("numcodecs.bitround", BitRound, qualname="zarr.codecs.numcodecs.BitRound")
+register_codec("numcodecs.blosc", Blosc, qualname="zarr.codecs.numcodecs.Blosc")
+register_codec("numcodecs.delta", Delta, qualname="zarr.codecs.numcodecs.Delta")
+register_codec(
+    "numcodecs.fixedscaleoffset",
+    FixedScaleOffset,
+    qualname="zarr.codecs.numcodecs.FixedScaleOffset",
+)
+register_codec("numcodecs.fletcher32", Fletcher32, qualname="zarr.codecs.numcodecs.Fletcher32")
+register_codec("numcodecs.gzip", GZip, qualname="zarr.codecs.numcodecs.GZip")
+register_codec(
+    "numcodecs.jenkins_lookup3", JenkinsLookup3, qualname="zarr.codecs.numcodecs.JenkinsLookup3"
+)
+register_codec("numcodecs.pcodec", PCodec, qualname="zarr.codecs.numcodecs.pcodec")
+register_codec("numcodecs.packbits", PackBits, qualname="zarr.codecs.numcodecs.PackBits")
+register_codec("numcodecs.quantize", Quantize, qualname="zarr.codecs.numcodecs.Quantize")
+register_codec("numcodecs.shuffle", Shuffle, qualname="zarr.codecs.numcodecs.Shuffle")
+register_codec("numcodecs.zlib", Zlib, qualname="zarr.codecs.numcodecs.Zlib")
+register_codec("numcodecs.zstd", Zstd, qualname="zarr.codecs.numcodecs.Zstd")

src/zarr/codecs/blosc.py

@@ -1,100 +1,59 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Mapping
 from dataclasses import dataclass, replace
+from enum import Enum
 from functools import cached_property
-from typing import (
-    TYPE_CHECKING,
-    Final,
-    Literal,
-    NotRequired,
-    TypedDict,
-    TypeGuard,
-    overload,
-)
+from typing import TYPE_CHECKING
 
 import numcodecs
 from numcodecs.blosc import Blosc
 from packaging.version import Version
-from typing_extensions import ReadOnly
-
-from zarr.abc.codec import BytesBytesCodec, CodecJSON
-from zarr.core.common import (
-    JSON,
-    NamedRequiredConfig,
-    ZarrFormat,
-)
+
+from zarr.abc.codec import BytesBytesCodec
+from zarr.core.buffer.cpu import as_numpy_array_wrapper
+from zarr.core.common import JSON, parse_enum, parse_named_configuration
 from zarr.core.dtype.common import HasItemSize
-from zarr.errors import CodecValidationError
-from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing import Self
 
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer
 
-BloscShuffle = Literal["noshuffle", "shuffle", "bitshuffle"]
-BLOSC_SHUFFLE: Final = ("noshuffle", "shuffle", "bitshuffle")
-
-BloscCname = Literal["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"]
-BLOSC_CNAME: Final = ("lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib")
-
-
-class BloscConfigV2(TypedDict):
-    cname: BloscCname
-    clevel: int
-    shuffle: int
-    blocksize: int
-    typesize: NotRequired[int]
-
 
-class BloscConfigV3(TypedDict):
-    cname: BloscCname
-    clevel: int
-    shuffle: BloscShuffle
-    blocksize: int
-    typesize: int
-
-
-class BloscJSON_V2(BloscConfigV2):
+class BloscShuffle(Enum):
     """
-    The JSON form of the Blosc codec in Zarr V2.
+    Enum for shuffle filter used by blosc.
     """
 
-    id: ReadOnly[Literal["blosc"]]
+    noshuffle = "noshuffle"
+    shuffle = "shuffle"
+    bitshuffle = "bitshuffle"
 
+    @classmethod
+    def from_int(cls, num: int) -> BloscShuffle:
+        blosc_shuffle_int_to_str = {
+            0: "noshuffle",
+            1: "shuffle",
+            2: "bitshuffle",
+        }
+        if num not in blosc_shuffle_int_to_str:
+            raise ValueError(f"Value must be between 0 and 2. Got {num}.")
+        return BloscShuffle[blosc_shuffle_int_to_str[num]]
 
-class BloscJSON_V3(NamedRequiredConfig[Literal["blosc"], BloscConfigV3]):
+
+class BloscCname(Enum):
     """
-    The JSON form of the Blosc codec in Zarr V3.
+    Enum for compression library used by blosc.
     """
 
-
-def check_json_v2(data: CodecJSON) -> TypeGuard[BloscJSON_V2]:
-    return (
-        isinstance(data, Mapping)
-        and set(data.keys()) == {"id", "clevel", "cname", "shuffle", "blocksize"}
-        and data["id"] == "blosc"
-    )
-
-
-def check_json_v3(data: CodecJSON) -> TypeGuard[BloscJSON_V3]:
-    return (
-        isinstance(data, Mapping)
-        and set(data.keys()) == {"name", "configuration"}
-        and data["name"] == "blosc"
-        and isinstance(data["configuration"], Mapping)
-        and set(data["configuration"].keys())
-        == {"cname", "clevel", "shuffle", "blocksize", "typesize"}
-    )
-
-
-def parse_cname(value: object) -> BloscCname:
-    if value not in BLOSC_CNAME:
-        raise ValueError(f"Value must be one of {BLOSC_CNAME}. Got {value} instead.")
-    return value
+    lz4 = "lz4"
+    lz4hc = "lz4hc"
+    blosclz = "blosclz"
+    zstd = "zstd"
+    snappy = "snappy"
+    zlib = "zlib"
 
 
 # See https://zarr.readthedocs.io/en/stable/user-guide/performance.html#configuring-blosc
@@ -125,35 +84,31 @@ def parse_blocksize(data: JSON) -> int:
     raise TypeError(f"Value should be an int. Got {type(data)} instead.")
 
 
-def parse_shuffle(data: object) -> BloscShuffle:
-    if data in BLOSC_SHUFFLE:
-        return data  # type: ignore[return-value]
-    raise TypeError(f"Value must be one of {BLOSC_SHUFFLE}. Got {data} instead.")
-
-
 @dataclass(frozen=True)
 class BloscCodec(BytesBytesCodec):
+    """blosc codec"""
+
     is_fixed_size = False
 
     typesize: int | None
-    cname: BloscCname
-    clevel: int
-    shuffle: BloscShuffle | None
-    blocksize: int
+    cname: BloscCname = BloscCname.zstd
+    clevel: int = 5
+    shuffle: BloscShuffle | None = BloscShuffle.noshuffle
+    blocksize: int = 0
 
     def __init__(
         self,
         *,
         typesize: int | None = None,
-        cname: BloscCname = "zstd",
+        cname: BloscCname | str = BloscCname.zstd,
         clevel: int = 5,
-        shuffle: BloscShuffle | None = None,
+        shuffle: BloscShuffle | str | None = None,
         blocksize: int = 0,
     ) -> None:
         typesize_parsed = parse_typesize(typesize) if typesize is not None else None
-        cname_parsed = parse_cname(cname)
+        cname_parsed = parse_enum(cname, BloscCname)
         clevel_parsed = parse_clevel(clevel)
-        shuffle_parsed = parse_shuffle(shuffle) if shuffle is not None else None
+        shuffle_parsed = parse_enum(shuffle, BloscShuffle) if shuffle is not None else None
         blocksize_parsed = parse_blocksize(blocksize)
 
         object.__setattr__(self, "typesize", typesize_parsed)
@@ -164,74 +119,24 @@ def __init__(
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self:
-        return cls.from_json(data, zarr_format=3)
+        _, configuration_parsed = parse_named_configuration(data, "blosc")
+        return cls(**configuration_parsed)  # type: ignore[arg-type]
 
     def to_dict(self) -> dict[str, JSON]:
-        return self.to_json(zarr_format=3)
-
-    @classmethod
-    def _from_json_v2(cls, data: CodecJSON) -> Self:
-        if check_json_v2(data):
-            return cls(
-                cname=data["cname"],
-                clevel=data["clevel"],
-                shuffle=BLOSC_SHUFFLE[data["shuffle"]],
-                blocksize=data["blocksize"],
-                typesize=data.get("typesize", None),
-            )
-        msg = (
-            "Invalid Zarr V2 JSON representation of the blosc codec. "
-            f"Got {data!r}, expected a Mapping with keys ('id', 'cname', 'clevel', 'shuffle', 'blocksize', 'typesize')"
-        )
-        raise CodecValidationError(msg)
-
-    @classmethod
-    def _from_json_v3(cls, data: CodecJSON) -> Self:
-        if check_json_v3(data):
-            return cls(
-                typesize=data["configuration"]["typesize"],
-                cname=data["configuration"]["cname"],
-                clevel=data["configuration"]["clevel"],
-                shuffle=data["configuration"]["shuffle"],
-                blocksize=data["configuration"]["blocksize"],
-            )
-        msg = (
-            "Invalid Zarr V3 JSON representation of the blosc codec. "
-            f"Got {data!r}, expected a Mapping with keys ('name', 'configuration')"
-            "Where the 'configuration' key is a Mapping with keys ('cname', 'clevel', 'shuffle', 'blocksize', 'typesize')"
-        )
-        raise CodecValidationError(msg)
-
-    @overload
-    def to_json(self, zarr_format: Literal[2]) -> BloscJSON_V2: ...
-    @overload
-    def to_json(self, zarr_format: Literal[3]) -> BloscJSON_V3: ...
-
-    def to_json(self, zarr_format: ZarrFormat) -> BloscJSON_V2 | BloscJSON_V3:
-        if self.typesize is None or self.shuffle is None:
-            raise ValueError("typesize and blocksize need to be set for encoding.")
-        if zarr_format == 2:
-            return {
-                "id": "blosc",
+        if self.typesize is None:
+            raise ValueError("`typesize` needs to be set for serialization.")
+        if self.shuffle is None:
+            raise ValueError("`shuffle` needs to be set for serialization.")
+        return {
+            "name": "blosc",
+            "configuration": {
+                "typesize": self.typesize,
+                "cname": self.cname.value,
                 "clevel": self.clevel,
-                "cname": self.cname,
-                "shuffle": BLOSC_SHUFFLE.index(self.shuffle),
+                "shuffle": self.shuffle.value,
                 "blocksize": self.blocksize,
-            }
-        elif zarr_format == 3:
-            return {
-                "name": "blosc",
-                "configuration": {
-                    "clevel": self.clevel,
-                    "cname": self.cname,
-                    "shuffle": self.shuffle,
-                    "typesize": self.typesize,
-                    "blocksize": self.blocksize,
-                },
-            }
-        raise ValueError(
-            f"Unsupported Zarr format {zarr_format}. Expected 2 or 3."
-        )  # pragma: no cover
+            },
+        }
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         item_size = 1
@@ -241,18 +146,26 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         if new_codec.typesize is None:
             new_codec = replace(new_codec, typesize=item_size)
         if new_codec.shuffle is None:
-            new_codec = replace(new_codec, shuffle="bitshuffle" if item_size == 1 else "shuffle")
+            new_codec = replace(
+                new_codec,
+                shuffle=(BloscShuffle.bitshuffle if item_size == 1 else BloscShuffle.shuffle),
+            )
 
         return new_codec
 
     @cached_property
     def _blosc_codec(self) -> Blosc:
         if self.shuffle is None:
             raise ValueError("`shuffle` needs to be set for decoding and encoding.")
+        map_shuffle_str_to_int = {
+            BloscShuffle.noshuffle: 0,
+            BloscShuffle.shuffle: 1,
+            BloscShuffle.bitshuffle: 2,
+        }
         config_dict = {
-            "cname": self.cname,
+            "cname": self.cname.name,
             "clevel": self.clevel,
-            "shuffle": BLOSC_SHUFFLE.index(self.shuffle),
+            "shuffle": map_shuffle_str_to_int[self.shuffle],
             "blocksize": self.blocksize,
         }
         # See https://github.com/zarr-developers/numcodecs/pull/713
@@ -265,8 +178,6 @@ async def _decode_single(
         chunk_bytes: Buffer,
         chunk_spec: ArraySpec,
     ) -> Buffer:
-        from zarr.core.buffer.cpu import as_numpy_array_wrapper
-
         return await asyncio.to_thread(
             as_numpy_array_wrapper, self._blosc_codec.decode, chunk_bytes, chunk_spec.prototype
         )
@@ -287,6 +198,3 @@ async def _encode_single(
 
     def compute_encoded_size(self, _input_byte_length: int, _chunk_spec: ArraySpec) -> int:
         raise NotImplementedError
-
-
-register_codec("blosc", BloscCodec)