dtype adaptor for zarr 3.1

Author: d-v-b

numcodecs/zarr3.py

@@ -33,13 +33,14 @@
 from warnings import warn
 
 import numpy as np
-
+from packaging.version import Version
+from importlib.metadata import version
 import numcodecs
 
 try:
     import zarr
 
-    if zarr.__version__ < "3.0.0":  # pragma: no cover
+    if version('zarr') < Version("3.0.0"):  # pragma: no cover
         raise ImportError("zarr 3.0.0 or later is required to use the numcodecs zarr integration.")
 except ImportError as e:  # pragma: no cover
     raise ImportError(
@@ -55,6 +56,19 @@
 
 CODEC_PREFIX = "numcodecs."
 
+def from_zarr_dtype(dtype: Any) -> np.dtype:
+    """
+    Get a numpy data type from an array spec, depending on the zarr version.
+    """
+    if version('zarr') >= Version("3.1.0"):
+        return dtype.to_native_dtype()
+    return dtype
+
+def to_zarr_dtype(dtype: np.dtype) -> Any:
+    if version('zarr') >= Version("3.1.0"):
+        from zarr.dtype import parse_data_type
+        return parse_data_type(dtype)
+    return dtype
 
 def _expect_name_prefix(codec_name: str) -> str:
     if not codec_name.startswith(CODEC_PREFIX):
@@ -224,15 +238,17 @@ class LZMA(_NumcodecsBytesBytesCodec, codec_name="lzma"):
 class Shuffle(_NumcodecsBytesBytesCodec, codec_name="shuffle"):
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Shuffle:
         if self.codec_config.get("elementsize") is None:
-            return Shuffle(**{**self.codec_config, "elementsize": array_spec.dtype.itemsize})
+            dtype = from_zarr_dtype(array_spec.dtype)
+            return Shuffle(**{**self.codec_config, "elementsize": dtype.itemsize})
         return self  # pragma: no cover
 
 
 # array-to-array codecs ("filters")
 class Delta(_NumcodecsArrayArrayCodec, codec_name="delta"):
     def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
         if astype := self.codec_config.get("astype"):
-            return replace(chunk_spec, dtype=np.dtype(astype))  # type: ignore[call-overload]
+            dtype = to_zarr_dtype(np.dtype(astype))
+            return replace(chunk_spec, dtype=dtype)  # type: ignore[call-overload]
         return chunk_spec
 
 
@@ -243,12 +259,14 @@ class BitRound(_NumcodecsArrayArrayCodec, codec_name="bitround"):
 class FixedScaleOffset(_NumcodecsArrayArrayCodec, codec_name="fixedscaleoffset"):
     def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
         if astype := self.codec_config.get("astype"):
-            return replace(chunk_spec, dtype=np.dtype(astype))  # type: ignore[call-overload]
+            dtype = to_zarr_dtype(np.dtype(astype))
+            return replace(chunk_spec, dtype=dtype)  # type: ignore[call-overload]
         return chunk_spec
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> FixedScaleOffset:
         if self.codec_config.get("dtype") is None:
-            return FixedScaleOffset(**{**self.codec_config, "dtype": str(array_spec.dtype)})
+            dtype = from_zarr_dtype(array_spec.dtype)
+            return FixedScaleOffset(**{**self.codec_config, "dtype": str(dtype)})
         return self
 
 
@@ -258,7 +276,8 @@ def __init__(self, **codec_config: JSON) -> None:
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Quantize:
         if self.codec_config.get("dtype") is None:
-            return Quantize(**{**self.codec_config, "dtype": str(array_spec.dtype)})
+            dtype = from_zarr_dtype(array_spec.dtype)
+            return Quantize(**{**self.codec_config, "dtype": str(dtype)})
         return self
 
 
@@ -271,17 +290,20 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
         )
 
     def validate(self, *, dtype: np.dtype[Any], **_kwargs) -> None:
-        if dtype != np.dtype("bool"):
+        _dtype = from_zarr_dtype(dtype)
+        if _dtype != np.dtype("bool"):
             raise ValueError(f"Packbits filter requires bool dtype. Got {dtype}.")
 
 
 class AsType(_NumcodecsArrayArrayCodec, codec_name="astype"):
     def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec:
-        return replace(chunk_spec, dtype=np.dtype(self.codec_config["encode_dtype"]))  # type: ignore[arg-type]
+        dtype = to_zarr_dtype(np.dtype(self.codec_config["encode_dtype"]))
+        return replace(chunk_spec, dtype=dtype)  # type: ignore[arg-type]
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> AsType:
         if self.codec_config.get("decode_dtype") is None:
-            return AsType(**{**self.codec_config, "decode_dtype": str(array_spec.dtype)})
+            dtype = from_zarr_dtype(array_spec.dtype)
+            return AsType(**{**self.codec_config, "decode_dtype": str(dtype)})
         return self