refactor: use inheritance to remove boilerplate in dtype definitions

Author: d-v-b

src/zarr/core/dtype/_numpy.py

@@ -31,7 +31,7 @@ def check_json_bool(data: JSON) -> TypeGuard[bool]:
     Bool
         True if the data is a boolean, False otherwise.
     """
-    return bool(isinstance(data, bool))
+    return isinstance(data, bool)
 
 
 def check_json_str(data: JSON) -> TypeGuard[str]:
@@ -293,7 +293,7 @@ def complex_to_json_v3(data: complex | np.complexfloating[Any, Any]) -> tuple[JS
     return float_to_json_v3(data.real), float_to_json_v3(data.imag)
 
 
-def complex_to_json(
+def complex_float_to_json(
     data: complex | np.complexfloating[Any, Any], zarr_format: ZarrFormat
 ) -> tuple[JSONFloat, JSONFloat]:
     """
@@ -424,60 +424,48 @@ def float_from_json(data: JSONFloat, zarr_format: ZarrFormat) -> float:
         return float_from_json_v3(data)
 
 
-def complex_from_json_v2(
-    data: tuple[JSONFloat, JSONFloat], dtype: np.dtypes.Complex64DType | np.dtypes.Complex128DType
-) -> np.complexfloating[Any, Any]:
+def complex_float_from_json_v2(data: tuple[JSONFloat, JSONFloat]) -> complex:
     """
     Convert a JSON complex float to a complex number (v2).
 
     Parameters
     ----------
     data : tuple[JSONFloat, JSONFloat]
         The JSON complex float to convert.
-    dtype : Any
-        The numpy dtype.
 
     Returns
     -------
     np.complexfloating
         The complex number.
     """
-    return dtype.type(complex(float_from_json_v2(data[0]), float_from_json_v2(data[1])))
+    return complex(float_from_json_v2(data[0]), float_from_json_v2(data[1]))
 
 
-def complex_from_json_v3(
-    data: tuple[JSONFloat, JSONFloat], dtype: np.dtypes.Complex64DType | np.dtypes.Complex128DType
-) -> np.complexfloating[Any, Any]:
+def complex_float_from_json_v3(data: tuple[JSONFloat, JSONFloat]) -> complex:
     """
     Convert a JSON complex float to a complex number (v3).
 
     Parameters
     ----------
     data : tuple[JSONFloat, JSONFloat]
         The JSON complex float to convert.
-    dtype : Any
-        The numpy dtype.
 
     Returns
     -------
     np.complexfloating
         The complex number.
     """
-    return dtype.type(complex(float_from_json_v3(data[0]), float_from_json_v3(data[1])))
+    return complex(float_from_json_v3(data[0]), float_from_json_v3(data[1]))
 
 
-def complex_from_json(
-    data: tuple[JSONFloat, JSONFloat], dtype: Any, zarr_format: ZarrFormat
-) -> np.complexfloating[Any, Any]:
+def complex_float_from_json(data: tuple[JSONFloat, JSONFloat], zarr_format: ZarrFormat) -> complex:
     """
     Convert a JSON complex float to a complex number based on zarr format.
 
     Parameters
     ----------
     data : tuple[JSONFloat, JSONFloat]
         The JSON complex float to convert.
-    dtype : Any
-        The numpy dtype.
     zarr_format : ZarrFormat
         The zarr format version.
 
@@ -487,12 +475,9 @@ def complex_from_json(
         The complex number.
     """
     if zarr_format == 2:
-        return complex_from_json_v2(data, dtype)
+        return complex_float_from_json_v2(data)
     else:
-        if check_json_complex_float_v3(data):
-            return complex_from_json_v3(data, dtype)
-        else:
-            raise TypeError(f"Invalid type: {data}. Expected a sequence of two numbers.")
+        return complex_float_from_json_v3(data)
     raise ValueError(f"Invalid zarr format: {zarr_format}. Expected 2 or 3.")
 
 

src/zarr/core/dtype/common.py

@@ -17,13 +17,14 @@
 # This is the bound for the dtypes that we support. If we support non-numpy dtypes,
 # then this bound will need to be widened.
 _BaseDType = np.dtype[np.generic]
-TScalar = TypeVar("TScalar", bound=_BaseScalar)
+TScalar_co = TypeVar("TScalar_co", bound=_BaseScalar, covariant=True)
 # TODO: figure out an interface or protocol that non-numpy dtypes can use
-TDType = TypeVar("TDType", bound=_BaseDType)
+# These two type parameters are covariant because we want isinstance(ZDType[Subclass](), ZDType[BaseDType]) to be True
+TDType_co = TypeVar("TDType_co", bound=_BaseDType, covariant=True)
 
 
 @dataclass(frozen=True, kw_only=True, slots=True)
-class ZDType(Generic[TDType, TScalar], ABC):
+class ZDType(Generic[TDType_co, TScalar_co], ABC):
     """
     Abstract base class for wrapping native array data types, e.g. numpy dtypes
 
@@ -41,11 +42,11 @@ class ZDType(Generic[TDType, TScalar], ABC):
     # mypy currently disallows class variables to contain type parameters
     # but it seems OK for us to use it here:
     # https://github.com/python/typing/discussions/1424#discussioncomment-7989934
-    dtype_cls: ClassVar[type[TDType]]  # type: ignore[misc]
+    dtype_cls: ClassVar[type[TDType_co]]  # type: ignore[misc]
     _zarr_v3_name: ClassVar[str]
 
     @classmethod
-    def check_dtype(cls: type[Self], dtype: _BaseDType) -> TypeGuard[TDType]:
+    def check_dtype(cls: type[Self], dtype: _BaseDType) -> TypeGuard[TDType_co]:
         """
         Check that a data type matches the dtype_cls class attribute. Used as a type guard.
 
@@ -89,7 +90,7 @@ def from_dtype(cls: type[Self], dtype: _BaseDType) -> Self:
 
     @classmethod
     @abstractmethod
-    def _from_dtype_unsafe(cls: type[Self], dtype: TDType) -> Self:
+    def _from_dtype_unsafe(cls: type[Self], dtype: _BaseDType) -> Self:
         """
         Wrap a native dtype without checking.
 
@@ -106,7 +107,7 @@ def _from_dtype_unsafe(cls: type[Self], dtype: TDType) -> Self:
         ...
 
     @abstractmethod
-    def to_dtype(self: Self) -> TDType:
+    def to_dtype(self: Self) -> TDType_co:
         """
         Return an instance of the wrapped dtype.
 
@@ -117,8 +118,61 @@ def to_dtype(self: Self) -> TDType:
         """
         ...
 
+    def cast_value(self, data: object) -> TScalar_co:
+        """
+        Cast a value to the wrapped scalar type. The type is first checked for compatibility. If it's
+        incompatible with the associated scalar type, a ``TypeError`` will be raised.
+
+        Parameters
+        ----------
+        data : TScalar
+            The scalar value to cast.
+
+        Returns
+        -------
+        TScalar
+            The cast value.
+        """
+        if self.check_value(data):
+            return self._cast_value_unsafe(data)
+        raise TypeError(f"Invalid value: {data}")
+
+    @abstractmethod
+    def check_value(self, data: object) -> bool:
+        """
+        Check that a value is a valid value for the wrapped data type.
+
+        Parameters
+        ----------
+        data : object
+            A value to check.
+
+        Returns
+        -------
+        Bool
+            True if the value is valid, False otherwise.
+        """
+        ...
+
+    @abstractmethod
+    def _cast_value_unsafe(self, data: object) -> TScalar_co:
+        """
+        Cast a value to the wrapped data type. This method should not perform any input validation.
+
+        Parameters
+        ----------
+        data : TScalar
+            The scalar value to cast.
+
+        Returns
+        -------
+        TScalar
+            The cast value.
+        """
+        ...
+
     @abstractmethod
-    def default_value(self) -> TScalar:
+    def default_value(self) -> TScalar_co:
         """
         Get the default value for the wrapped data type. This is a method, rather than an attribute,
         because the default value for some data types may depend on parameters that are not known
@@ -216,7 +270,7 @@ def _from_json_unsafe(cls: type[Self], data: JSON, zarr_format: ZarrFormat) -> S
         ...
 
     @abstractmethod
-    def to_json_value(self, data: TScalar, *, zarr_format: ZarrFormat) -> JSON:
+    def to_json_value(self, data: object, *, zarr_format: ZarrFormat) -> JSON:
         """
         Convert a single value to JSON-serializable format.
 
@@ -235,7 +289,7 @@ def to_json_value(self, data: TScalar, *, zarr_format: ZarrFormat) -> JSON:
         ...
 
     @abstractmethod
-    def from_json_value(self: Self, data: JSON, *, zarr_format: ZarrFormat) -> TScalar:
+    def from_json_value(self: Self, data: JSON, *, zarr_format: ZarrFormat) -> TScalar_co:
         """
         Read a JSON-serializable value as a scalar.
 

src/zarr/core/dtype/wrapper.py

@@ -9,7 +9,7 @@
 
 from zarr.abc.metadata import Metadata
 from zarr.core.dtype import get_data_type_from_native_dtype
-from zarr.core.dtype.wrapper import TDType, TScalar, ZDType, _BaseDType, _BaseScalar
+from zarr.core.dtype.wrapper import TDType_co, TScalar_co, ZDType, _BaseDType, _BaseScalar
 
 if TYPE_CHECKING:
     from typing import Any, Literal, Self
@@ -58,7 +58,7 @@ def __init__(
         self,
         *,
         shape: ChunkCoords,
-        dtype: ZDType[TDType, TScalar],
+        dtype: ZDType[TDType_co, TScalar_co],
         chunks: ChunkCoords,
         fill_value: Any,
         order: MemoryOrder,
@@ -176,7 +176,7 @@ def to_dict(self) -> dict[str, JSON]:
             zarray_dict["filters"] = new_filters
 
         if self.fill_value is not None:
-            fill_value = self.dtype.to_json_value(self.fill_value, zarr_format=2)  # type: ignore[arg-type]
+            fill_value = self.dtype.to_json_value(self.fill_value, zarr_format=2)
             zarray_dict["fill_value"] = fill_value
 
         zarray_dict["dtype"] = self.dtype.to_json(zarr_format=2)

src/zarr/core/metadata/v2.py

@@ -20,7 +20,8 @@
 from zarr.core.chunk_grids import RegularChunkGrid, _auto_partition
 from zarr.core.common import JSON, parse_shapelike
 from zarr.core.config import config as zarr_config
-from zarr.core.dtype import get_data_type_from_native_dtype
+from zarr.core.dtype import data_type_registry, get_data_type_from_native_dtype
+from zarr.core.dtype._numpy import DateTime64, HasLength, Structured
 from zarr.core.metadata.v2 import ArrayV2Metadata
 from zarr.core.metadata.v3 import ArrayV3Metadata
 from zarr.core.sync import sync
@@ -36,6 +37,7 @@
     from zarr.core.array import CompressorsLike, FiltersLike, SerializerLike, ShardsLike
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding, ChunkKeyEncodingLike
     from zarr.core.common import ChunkCoords, MemoryOrder, ShapeLike, ZarrFormat
+    from zarr.core.dtype.wrapper import ZDType
 
 
 async def parse_store(
@@ -404,3 +406,17 @@ def meta_from_array(
         chunk_key_encoding=chunk_key_encoding,
         dimension_names=dimension_names,
     )
+
+
+# Generate a collection of zdtype instances for use in testing.
+zdtype_examples: tuple[ZDType[Any, Any], ...] = ()
+for wrapper_cls in data_type_registry.contents.values():
+    # The Structured dtype has to be constructed with some actual fields
+    if wrapper_cls is Structured:
+        zdtype_examples += (wrapper_cls.from_dtype(np.dtype([("a", np.float64), ("b", np.int8)])),)
+    elif issubclass(wrapper_cls, HasLength):
+        zdtype_examples += (wrapper_cls(length=1),)
+    elif issubclass(wrapper_cls, DateTime64):
+        zdtype_examples += (wrapper_cls(unit="s"),)
+    else:
+        zdtype_examples += (wrapper_cls(),)

tests/conftest.py

@@ -40,12 +40,15 @@
 from zarr.core.common import JSON, MemoryOrder, ZarrFormat
 from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.dtype._numpy import Float64
+from zarr.core.dtype.wrapper import ZDType
 from zarr.core.group import AsyncGroup
 from zarr.core.indexing import BasicIndexer, ceildiv
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
 from zarr.storage import LocalStore, MemoryStore, StorePath
 
+from .conftest import zdtype_examples
+
 if TYPE_CHECKING:
     from zarr.core.array_spec import ArrayConfigLike
     from zarr.core.metadata.v2 import ArrayV2Metadata
@@ -177,32 +180,42 @@ def test_array_name_properties_with_group(
 
 @pytest.mark.parametrize("store", ["memory"], indirect=True)
 @pytest.mark.parametrize("specifiy_fill_value", [True, False])
-@pytest.mark.parametrize("dtype_str", ["bool", "uint8", "complex64"])
-def test_array_v3_fill_value_default(
-    store: MemoryStore, specifiy_fill_value: bool, dtype_str: str
+@pytest.mark.parametrize(
+    "zdtype", zdtype_examples, ids=tuple(str(type(v)) for v in zdtype_examples)
+)
+def test_array_fill_value_default(
+    store: MemoryStore, specifiy_fill_value: bool, zdtype: ZDType[Any, Any]
 ) -> None:
     """
     Test that creating an array with the fill_value parameter set to None, or unspecified,
     results in the expected fill_value attribute of the array, i.e. 0 cast to the array's dtype.
     """
     shape = (10,)
-    default_fill_value = 0
     if specifiy_fill_value:
         arr = zarr.create_array(
             store=store,
             shape=shape,
-            dtype=dtype_str,
+            dtype=zdtype,
             zarr_format=3,
             chunks=shape,
             fill_value=None,
         )
     else:
-        arr = zarr.create_array(
-            store=store, shape=shape, dtype=dtype_str, zarr_format=3, chunks=shape
-        )
+        arr = zarr.create_array(store=store, shape=shape, dtype=zdtype, zarr_format=3, chunks=shape)
+    expected_fill_value = zdtype.default_value()
+    if isinstance(expected_fill_value, np.datetime64 | np.timedelta64):
+        if np.isnat(expected_fill_value):
+            assert np.isnat(arr.fill_value)
+    elif isinstance(expected_fill_value, np.floating | np.complexfloating):
+        if np.isnan(expected_fill_value):
+            assert np.isnan(arr.fill_value)
+    else:
+        assert arr.fill_value == expected_fill_value
+    # A simpler check would be to ensure that arr.fill_value.dtype == arr.dtype
+    # But for some numpy data types (namely, U), scalars might not have length. An empty string
+    # scalar from a `>U4` array would have dtype `>U`, and arr.fill_value.dtype == arr.dtype will fail.
 
-    assert arr.fill_value == np.dtype(dtype_str).type(default_fill_value)
-    assert arr.fill_value.dtype == arr.dtype
+    assert type(arr.fill_value) is type(np.array([arr.fill_value], dtype=arr.dtype)[0])
 
 
 @pytest.mark.parametrize("store", ["memory"], indirect=True)
@@ -1004,7 +1017,7 @@ async def test_v3_chunk_encoding(
             filters=filters,
             compressors=compressors,
             serializer="auto",
-            dtype=arr.metadata.data_type,  # type: ignore[union-attr]
+            dtype=arr._zdtype,
         )
         assert arr.filters == filters_expected
         assert arr.compressors == compressors_expected
@@ -1369,4 +1382,4 @@ async def test_sharding_coordinate_selection() -> None:
         shards=(2, 4, 4),
     )
     arr[:] = np.arange(2 * 3 * 4).reshape((2, 3, 4))
-    assert (arr[1, [0, 1]] == np.array([[12, 13, 14, 15], [16, 17, 18, 19]])).all()  # type: ignore[index]
+    assert (arr[1, [0, 1]] == np.array([[12, 13, 14, 15], [16, 17, 18, 19]])).all()