use a common function signature for from_json by packing the object_codec_id in a typeddict for zarr v2 metadata

Author: d-v-b

src/zarr/core/dtype/__init__.py

@@ -2,7 +2,10 @@
 
 from typing import TYPE_CHECKING, Final, TypeAlias
 
-from zarr.core.dtype.common import DataTypeValidationError
+from zarr.core.dtype.common import (
+    DataTypeValidationError,
+    DTypeJSON,
+)
 from zarr.core.dtype.npy.bool import Bool
 from zarr.core.dtype.npy.bytes import NullTerminatedBytes, RawBytes, VariableLengthBytes
 from zarr.core.dtype.npy.complex import Complex64, Complex128
@@ -131,20 +134,20 @@ def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[TBaseDType,
     return data_type_registry.match_dtype(dtype=na_dtype)
 
 
-def get_data_type_from_json_v3(
-    dtype_spec: JSON,
-) -> ZDType[TBaseDType, TBaseScalar]:
-    return data_type_registry.match_json_v3(dtype_spec)
-
-
-def get_data_type_from_json_v2(
-    dtype_spec: JSON, *, object_codec_id: str | None = None
+def get_data_type_from_json(
+    dtype_spec: DTypeJSON, *, zarr_format: ZarrFormat
 ) -> ZDType[TBaseDType, TBaseScalar]:
-    return data_type_registry.match_json_v2(dtype_spec, object_codec_id=object_codec_id)
+    """
+    Given a JSON representation of a data type and a Zarr format version,
+    attempt to create a ZDType instance from the registered ZDType classes.
+    """
+    return data_type_registry.match_json(dtype_spec, zarr_format=zarr_format)
 
 
 def parse_data_type(
-    dtype_spec: ZDTypeLike, *, zarr_format: ZarrFormat, object_codec_id: str | None = None
+    dtype_spec: ZDTypeLike,
+    *,
+    zarr_format: ZarrFormat,
 ) -> ZDType[TBaseDType, TBaseScalar]:
     """
     Interpret the input as a ZDType instance.
@@ -153,7 +156,7 @@ def parse_data_type(
         return dtype_spec
     # dict and zarr_format 3 means that we have a JSON object representation of the dtype
     if zarr_format == 3 and isinstance(dtype_spec, Mapping):
-        return get_data_type_from_json_v3(dtype_spec)  # type: ignore[arg-type]
+        return get_data_type_from_json(dtype_spec, zarr_format=3)
     # otherwise, we have either a numpy dtype string, or a zarr v3 dtype string, and in either case
     # we can create a numpy dtype from it, and do the dtype inference from that
     return get_data_type_from_native_dtype(dtype_spec)  # type: ignore[arg-type]

src/zarr/core/dtype/common.py

@@ -1,16 +1,149 @@
 from __future__ import annotations
 
 import warnings
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
-from typing import ClassVar, Final, Literal
+from typing import (
+    ClassVar,
+    Final,
+    Generic,
+    Literal,
+    TypedDict,
+    TypeGuard,
+    TypeVar,
+)
+
+from zarr.core.common import NamedConfig
 
 EndiannessStr = Literal["little", "big"]
 ENDIANNESS_STR: Final = "little", "big"
+
 SpecialFloatStrings = Literal["NaN", "Infinity", "-Infinity"]
 SPECIAL_FLOAT_STRINGS: Final = ("NaN", "Infinity", "-Infinity")
+
 JSONFloatV2 = float | SpecialFloatStrings
 JSONFloatV3 = float | SpecialFloatStrings | str
 
+ObjectCodecID = Literal["vlen-utf8", "vlen-bytes", "vlen-array", "pickle", "json2", "msgpack2"]
+# These are the ids of the known object codecs for zarr v2.
+OBJECT_CODEC_IDS: Final = ("vlen-utf8", "vlen-bytes", "vlen-array", "pickle", "json2", "msgpack2")
+
+# This is a wider type than our standard JSON type because we need
+# to work with typeddict objects which are assignable to Mapping[str, object]
+DTypeJSON = str | int | float | Sequence["DTypeJSON"] | None | Mapping[str, object]
+
+# The DTypeJSON_V2 type exists because ZDType.from_json takes a single argument, which must contain
+# all the information necessary to decode the data type. Zarr v2 supports multiple distinct
+# data types that all used the "|O" data type identifier. These data types can only be
+# discriminated on the basis of their "object codec", i.e. a special data type specific
+# compressor or filter. So to figure out what data type a zarr v2 array has, we need the
+# data type identifier from metadata, as well as an object codec id if the data type identifier
+# is "|O".
+# So we will pack the name of the dtype alongside the name of the object codec id, if applicable,
+# in a single dict, and pass that to the data type inference logic.
+# These type variables have a very wide bound because the individual zdtype
+# classes can perform a very specific type check.
+
+# This is the JSON representation of a structured dtype in zarr v2
+StructuredName_V2 = Sequence["str | StructuredName_V2"]
+
+# This models the type of the name a dtype might have in zarr v2 array metadata
+DTypeName_V2 = StructuredName_V2 | str
+
+TDTypeNameV2_co = TypeVar("TDTypeNameV2_co", bound=DTypeName_V2, covariant=True)
+TObjectCodecID_co = TypeVar("TObjectCodecID_co", bound=None | str, covariant=True)
+
+
+class DTypeConfig_V2(TypedDict, Generic[TDTypeNameV2_co, TObjectCodecID_co]):
+    name: TDTypeNameV2_co
+    object_codec_id: TObjectCodecID_co
+
+
+DTypeSpec_V2 = DTypeConfig_V2[DTypeName_V2, None | str]
+
+
+def check_structured_dtype_v2_inner(data: object) -> TypeGuard[StructuredName_V2]:
+    """
+    A type guard for the inner elements of a structured dtype. This is a recursive check because
+    the type is itself recursive.
+
+    This check ensures that all the elements are 2-element sequences beginning with a string
+    and ending with either another string or another 2-element sequence beginning with a string and
+    ending with another instance of that type.
+    """
+    if isinstance(data, (str, Mapping)):
+        return False
+    if not isinstance(data, Sequence):
+        return False
+    if len(data) != 2:
+        return False
+    if not (isinstance(data[0], str)):
+        return False
+    if isinstance(data[-1], str):
+        return True
+    elif isinstance(data[-1], Sequence):
+        return check_structured_dtype_v2_inner(data[-1])
+    return False
+
+
+def check_structured_dtype_name_v2(data: Sequence[object]) -> TypeGuard[StructuredName_V2]:
+    return all(check_structured_dtype_v2_inner(d) for d in data)
+
+
+def check_dtype_name_v2(data: object) -> TypeGuard[DTypeName_V2]:
+    """
+    Type guard for narrowing the type of a python object to an valid zarr v2 dtype name.
+    """
+    if isinstance(data, str):
+        return True
+    elif isinstance(data, Sequence):
+        return check_structured_dtype_name_v2(data)
+    return False
+
+
+def check_dtype_spec_v2(data: object) -> TypeGuard[DTypeSpec_V2]:
+    """
+    Type guard for narrowing a python object to an instance of DTypeSpec_V2
+    """
+    if not isinstance(data, Mapping):
+        return False
+    if set(data.keys()) != {"name", "object_codec_id"}:
+        return False
+    if not check_dtype_name_v2(data["name"]):
+        return False
+    return isinstance(data["object_codec_id"], str | None)
+
+
+# By comparison, The JSON representation of a dtype in zarr v3 is much simpler.
+# It's either a string, or a structured dict
+DTypeSpec_V3 = str | NamedConfig[str, Mapping[str, object]]
+
+
+def check_dtype_spec_v3(data: object) -> TypeGuard[DTypeSpec_V3]:
+    """
+    Type guard for narrowing the type of a python object to an instance of
+    DTypeSpec_V3, i.e either a string or a dict with a "name" field that's a string and a
+    "configuration" field that's a mapping with string keys.
+    """
+    if isinstance(data, str) or (  # noqa: SIM103
+        isinstance(data, Mapping)
+        and set(data.keys()) == {"name", "configuration"}
+        and isinstance(data["configuration"], Mapping)
+        and all(isinstance(k, str) for k in data["configuration"])
+    ):
+        return True
+    return False
+
+
+def unpack_dtype_json(data: DTypeSpec_V2 | DTypeSpec_V3) -> DTypeJSON:
+    """
+    Return the array metadata form of the dtype JSON representation. For the Zarr V3 form of dtype
+    metadata, this is a no-op. For the Zarr V2 form of dtype metadata, this unpacks the dtype name.
+    """
+    if isinstance(data, Mapping) and set(data.keys()) == {"name", "object_codec_id"}:
+        return data["name"]
+    return data
+
 
 class DataTypeValidationError(ValueError): ...
 

src/zarr/core/dtype/npy/bool.py

@@ -1,12 +1,22 @@
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import ClassVar, Literal, Self, TypeGuard, overload
+from typing import TYPE_CHECKING, ClassVar, Literal, Self, TypeGuard, overload
 
 import numpy as np
 
-from zarr.core.common import JSON, ZarrFormat
-from zarr.core.dtype.common import DataTypeValidationError, HasItemSize
+from zarr.core.dtype.common import (
+    DataTypeValidationError,
+    DTypeConfig_V2,
+    DTypeJSON,
+    HasItemSize,
+    check_dtype_spec_v2,
+)
 from zarr.core.dtype.wrapper import TBaseDType, ZDType
 
+if TYPE_CHECKING:
+    from zarr.core.common import JSON, ZarrFormat
+
 
 @dataclass(frozen=True, kw_only=True, slots=True)
 class Bool(ZDType[np.dtypes.BoolDType, np.bool_], HasItemSize):
@@ -44,40 +54,47 @@ def to_native_dtype(self: Self) -> np.dtypes.BoolDType:
 
     @classmethod
     def _check_json_v2(
-        cls, data: JSON, *, object_codec_id: str | None = None
-    ) -> TypeGuard[Literal["|b1"]]:
+        cls,
+        data: DTypeJSON,
+    ) -> TypeGuard[DTypeConfig_V2[Literal["|b1"], None]]:
         """
         Check that the input is a valid JSON representation of a Bool.
         """
-        return data == cls._zarr_v2_name
+        return (
+            check_dtype_spec_v2(data)
+            and data["name"] == cls._zarr_v2_name
+            and data["object_codec_id"] is None
+        )
 
     @classmethod
-    def _check_json_v3(cls, data: JSON) -> TypeGuard[Literal["bool"]]:
+    def _check_json_v3(cls, data: DTypeJSON) -> TypeGuard[Literal["bool"]]:
         return data == cls._zarr_v3_name
 
     @classmethod
-    def from_json_v2(cls, data: JSON, *, object_codec_id: str | None = None) -> Self:
+    def _from_json_v2(cls, data: DTypeJSON) -> Self:
         if cls._check_json_v2(data):
             return cls()
         msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v2_name!r}"
         raise DataTypeValidationError(msg)
 
     @classmethod
-    def from_json_v3(cls: type[Self], data: JSON) -> Self:
+    def _from_json_v3(cls: type[Self], data: DTypeJSON) -> Self:
         if cls._check_json_v3(data):
             return cls()
         msg = f"Invalid JSON representation of {cls.__name__}. Got {data!r}, expected the string {cls._zarr_v3_name!r}"
         raise DataTypeValidationError(msg)
 
-    @overload
-    def to_json(self, zarr_format: Literal[2]) -> Literal["|b1"]: ...
+    @overload  # type: ignore[override]
+    def to_json(self, zarr_format: Literal[2]) -> DTypeConfig_V2[Literal["|b1"], None]: ...
 
     @overload
     def to_json(self, zarr_format: Literal[3]) -> Literal["bool"]: ...
 
-    def to_json(self, zarr_format: ZarrFormat) -> Literal["|b1", "bool"]:
+    def to_json(
+        self, zarr_format: ZarrFormat
+    ) -> DTypeConfig_V2[Literal["|b1"], None] | Literal["bool"]:
         if zarr_format == 2:
-            return self._zarr_v2_name
+            return {"name": self._zarr_v2_name, "object_codec_id": None}
         elif zarr_format == 3:
             return self._zarr_v3_name
         raise ValueError(f"zarr_format must be 2 or 3, got {zarr_format}")  # pragma: no cover