add to_json methods to codecs

Author: d-v-b

src/zarr/codecs/blosc.py

@@ -1,18 +1,22 @@
 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
+from typing import TYPE_CHECKING, Final, Literal, NotRequired, TypedDict, TypeGuard, overload
 
 import numcodecs
 from numcodecs.blosc import Blosc
 from packaging.version import Version
 
-from zarr.abc.codec import BytesBytesCodec
+from zarr.abc.codec import BytesBytesCodec, CodecJSON, CodecJSON_V2, CodecValidationError
 from zarr.core.buffer.cpu import as_numpy_array_wrapper
-from zarr.core.common import JSON, parse_enum, parse_named_configuration
+from zarr.core.common import (
+    JSON,
+    NamedRequiredConfig,
+    ZarrFormat,
+)
 from zarr.core.dtype.common import HasItemSize
 from zarr.registry import register_codec
 
@@ -22,39 +26,64 @@
     from zarr.core.array_spec import ArraySpec
     from zarr.core.buffer import Buffer
 
+BloscShuffle = Literal["noshuffle", "shuffle", "bitshuffle"]
+BLOSC_SHUFFLE: Final = ("noshuffle", "shuffle", "bitshuffle")
 
-class BloscShuffle(Enum):
-    """
-    Enum for shuffle filter used by blosc.
-    """
+BloscCname = Literal["lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib"]
+BLOSC_CNAME: Final = ("lz4", "lz4hc", "blosclz", "zstd", "snappy", "zlib")
 
-    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 BloscConfigV2(TypedDict):
+    cname: BloscCname
+    clevel: int
+    shuffle: int
+    blocksize: int
+    typesize: NotRequired[int]
 
 
-class BloscCname(Enum):
+class BloscConfigV3(TypedDict):
+    cname: BloscCname
+    clevel: int
+    shuffle: BloscShuffle
+    blocksize: int
+    typesize: int
+
+
+class BloscJSON_V2(CodecJSON_V2[Literal["blosc"]], BloscConfigV2):
     """
-    Enum for compression library used by blosc.
+    The JSON form of the Blosc codec in Zarr V2.
     """
 
-    lz4 = "lz4"
-    lz4hc = "lz4hc"
-    blosclz = "blosclz"
-    zstd = "zstd"
-    snappy = "snappy"
-    zlib = "zlib"
+
+class BloscJSON_V3(NamedRequiredConfig[Literal["blosc"], BloscConfigV3]):
+    """
+    The JSON form of the Blosc codec in Zarr V3.
+    """
+
+
+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
 
 
 # See https://zarr.readthedocs.io/en/stable/user-guide/performance.html#configuring-blosc
@@ -85,31 +114,35 @@ 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 = BloscCname.zstd
-    clevel: int = 5
-    shuffle: BloscShuffle | None = BloscShuffle.noshuffle
-    blocksize: int = 0
+    cname: BloscCname
+    clevel: int
+    shuffle: BloscShuffle | None
+    blocksize: int
 
     def __init__(
         self,
         *,
         typesize: int | None = None,
-        cname: BloscCname | str = BloscCname.zstd,
+        cname: BloscCname = "zstd",
         clevel: int = 5,
-        shuffle: BloscShuffle | str | None = None,
+        shuffle: BloscShuffle | None = None,
         blocksize: int = 0,
     ) -> None:
         typesize_parsed = parse_typesize(typesize) if typesize is not None else None
-        cname_parsed = parse_enum(cname, BloscCname)
+        cname_parsed = parse_cname(cname)
         clevel_parsed = parse_clevel(clevel)
-        shuffle_parsed = parse_enum(shuffle, BloscShuffle) if shuffle is not None else None
+        shuffle_parsed = parse_shuffle(shuffle) if shuffle is not None else None
         blocksize_parsed = parse_blocksize(blocksize)
 
         object.__setattr__(self, "typesize", typesize_parsed)
@@ -120,24 +153,74 @@ def __init__(
 
     @classmethod
     def from_dict(cls, data: dict[str, JSON]) -> Self:
-        _, configuration_parsed = parse_named_configuration(data, "blosc")
-        return cls(**configuration_parsed)  # type: ignore[arg-type]
+        return cls.from_json(data, zarr_format=3)
 
     def to_dict(self) -> dict[str, JSON]:
-        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,
+        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",
                 "clevel": self.clevel,
-                "shuffle": self.shuffle.value,
+                "cname": self.cname,
+                "shuffle": BLOSC_SHUFFLE.index(self.shuffle),
                 "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
@@ -147,26 +230,18 @@ 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=(BloscShuffle.bitshuffle if item_size == 1 else BloscShuffle.shuffle),
-            )
+            new_codec = replace(new_codec, shuffle="bitshuffle" if item_size == 1 else "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.name,
+            "cname": self.cname,
             "clevel": self.clevel,
-            "shuffle": map_shuffle_str_to_int[self.shuffle],
+            "shuffle": BLOSC_SHUFFLE.index(self.shuffle),
             "blocksize": self.blocksize,
         }
         # See https://github.com/zarr-developers/numcodecs/pull/713