use wrap / unwrap instead of to_dtype / from_dtype; push into v2 codebase

Author: d-v-b

src/zarr/api/asynchronous.py

@@ -32,6 +32,7 @@
 from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
 from zarr.core.metadata.v2 import _default_compressor, _default_filters
 from zarr.errors import NodeTypeValidationError
+from zarr.registry import get_data_type_from_numpy
 from zarr.storage._common import make_store_path
 
 if TYPE_CHECKING:
@@ -428,11 +429,12 @@ async def save_array(
     shape = arr.shape
     chunks = getattr(arr, "chunks", None)  # for array-likes with chunks attribute
     overwrite = kwargs.pop("overwrite", None) or _infer_overwrite(mode)
+    zarr_dtype = get_data_type_from_numpy(arr.dtype)
     new = await AsyncArray._create(
         store_path,
         zarr_format=zarr_format,
         shape=shape,
-        dtype=arr.dtype,
+        dtype=zarr_dtype,
         chunks=chunks,
         overwrite=overwrite,
         **kwargs,
@@ -978,15 +980,14 @@ async def create(
         _handle_zarr_version_or_format(zarr_version=zarr_version, zarr_format=zarr_format)
         or _default_zarr_format()
     )
-
+    dtype_wrapped = parse_dtype(dtype, zarr_format=zarr_format)
     if zarr_format == 2:
         if chunks is None:
             chunks = shape
-        dtype = parse_dtype(dtype, zarr_format=zarr_format)
         if not filters:
-            filters = _default_filters(dtype)
+            filters = _default_filters(dtype_wrapped)
         if not compressor:
-            compressor = _default_compressor(dtype)
+            compressor = _default_compressor(dtype_wrapped)
     elif zarr_format == 3 and chunk_shape is None:  # type: ignore[redundant-expr]
         if chunks is not None:
             chunk_shape = chunks
@@ -1051,7 +1052,7 @@ async def create(
         store_path,
         shape=shape,
         chunks=chunks,
-        dtype=dtype,
+        dtype=dtype_wrapped,
         compressor=compressor,
         fill_value=fill_value,
         overwrite=overwrite,

src/zarr/codecs/_v2.py

@@ -48,15 +48,15 @@ async def _decode_single(
         # segfaults and other bad things happening
         if chunk_spec.dtype != object:
             try:
-                chunk = chunk.view(chunk_spec.dtype)
+                chunk = chunk.view(chunk_spec.dtype.unwrap())
             except TypeError:
                 # this will happen if the dtype of the chunk
                 # does not match the dtype of the array spec i.g. if
                 # the dtype of the chunk_spec is a string dtype, but the chunk
                 # is an object array. In this case, we need to convert the object
                 # array to the correct dtype.
 
-                chunk = np.array(chunk).astype(chunk_spec.dtype)
+                chunk = np.array(chunk).astype(chunk_spec.dtype.unwrap())
 
         elif chunk.dtype != object:
             # If we end up here, someone must have hacked around with the filters.
@@ -80,7 +80,7 @@ async def _encode_single(
         chunk = chunk_array.as_ndarray_like()
 
         # ensure contiguous and correct order
-        chunk = chunk.astype(chunk_spec.dtype, order=chunk_spec.order, copy=False)
+        chunk = chunk.astype(chunk_spec.dtype.unwrap(), order=chunk_spec.order, copy=False)
 
         # apply filters
         if self.filters:

src/zarr/codecs/bytes.py

@@ -56,7 +56,7 @@ def to_dict(self) -> dict[str, JSON]:
             return {"name": "bytes", "configuration": {"endian": self.endian.value}}
 
     def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
-        if array_spec.dtype.itemsize == 0:
+        if array_spec.dtype.unwrap().itemsize == 0:
             if self.endian is not None:
                 return replace(self, endian=None)
         elif self.endian is None:

src/zarr/core/array.py

@@ -98,7 +98,7 @@
     ArrayV3MetadataDict,
     T_ArrayMetadata,
 )
-from zarr.core.metadata.dtype import DTypeWrapper
+from zarr.core.metadata.dtype import DTypeWrapper, VariableLengthString
 from zarr.core.metadata.v2 import (
     _default_compressor,
     _default_filters,
@@ -549,7 +549,7 @@ async def _create(
         *,
         # v2 and v3
         shape: ShapeLike,
-        dtype: npt.DTypeLike,
+        dtype: npt.DTypeLike[Any],
         zarr_format: ZarrFormat = 3,
         fill_value: Any | None = None,
         attributes: dict[str, JSON] | None = None,
@@ -578,18 +578,22 @@ async def _create(
         See :func:`AsyncArray.create` for more details.
         Deprecated in favor of :func:`zarr.api.asynchronous.create_array`.
         """
+        # TODO: delete this and be more strict about where parsing occurs
+        if not isinstance(dtype, DTypeWrapper):
+            dtype_parsed = get_data_type_from_numpy(np.dtype(dtype))
+        else:
+            dtype_parsed = dtype
         store_path = await make_store_path(store)
 
-        dtype_parsed = parse_dtype(dtype, zarr_format=zarr_format)
         shape = parse_shapelike(shape)
 
         if chunks is not None and chunk_shape is not None:
             raise ValueError("Only one of chunk_shape or chunks can be provided.")
 
         if chunks:
-            _chunks = normalize_chunks(chunks, shape, dtype_parsed.itemsize)
+            _chunks = normalize_chunks(chunks, shape, dtype_parsed.unwrap().itemsize)
         else:
-            _chunks = normalize_chunks(chunk_shape, shape, dtype_parsed.itemsize)
+            _chunks = normalize_chunks(chunk_shape, shape, dtype_parsed.unwrap().itemsize)
         config_parsed = parse_array_config(config)
 
         result: AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]
@@ -666,7 +670,7 @@ async def _create(
     @staticmethod
     def _create_metadata_v3(
         shape: ShapeLike,
-        dtype: np.dtype[Any],
+        dtype: DTypeWrapper[Any, Any],
         chunk_shape: ChunkCoords,
         fill_value: Any | None = None,
         chunk_key_encoding: ChunkKeyEncodingLike | None = None,
@@ -694,19 +698,16 @@ def _create_metadata_v3(
                 stacklevel=2,
             )
 
-        # resolve the numpy dtype into zarr v3 datatype
-        zarr_data_type = get_data_type_from_numpy(dtype)
-
         if fill_value is None:
             # v3 spec will not allow a null fill value
-            fill_value_parsed = zarr_data_type.default_value
+            fill_value_parsed = dtype.default_value
         else:
             fill_value_parsed = fill_value
 
         chunk_grid_parsed = RegularChunkGrid(chunk_shape=chunk_shape)
         return ArrayV3Metadata(
             shape=shape,
-            data_type=zarr_data_type,
+            data_type=dtype,
             chunk_grid=chunk_grid_parsed,
             chunk_key_encoding=chunk_key_encoding_parsed,
             fill_value=fill_value_parsed,
@@ -769,7 +770,7 @@ async def _create_v3(
     @staticmethod
     def _create_metadata_v2(
         shape: ChunkCoords,
-        dtype: np.dtype[Any],
+        dtype: DTypeWrapper[Any, Any],
         chunks: ChunkCoords,
         order: MemoryOrder,
         dimension_separator: Literal[".", "/"] | None = None,
@@ -781,10 +782,8 @@ def _create_metadata_v2(
         if dimension_separator is None:
             dimension_separator = "."
 
-        dtype = parse_dtype(dtype, zarr_format=2)
-
         # inject VLenUTF8 for str dtype if not already present
-        if np.issubdtype(dtype, np.str_):
+        if isinstance(dtype, VariableLengthString):
             filters = filters or []
             from numcodecs.vlen import VLenUTF8
 
@@ -793,7 +792,7 @@ def _create_metadata_v2(
 
         return ArrayV2Metadata(
             shape=shape,
-            dtype=np.dtype(dtype),
+            dtype=dtype,
             chunks=chunks,
             order=order,
             dimension_separator=dimension_separator,
@@ -2046,7 +2045,7 @@ def dtype(self) -> np.dtype[Any]:
         np.dtype
             The NumPy data type.
         """
-        return self._async_array.dtype
+        return self._async_array.dtype.unwrap()
 
     @property
     def attrs(self) -> Attributes:
@@ -3919,7 +3918,7 @@ async def init_array(
 
     from zarr.codecs.sharding import ShardingCodec, ShardingCodecIndexLocation
 
-    dtype_parsed = parse_dtype(dtype, zarr_format=zarr_format)
+    dtype_wrapped = parse_dtype(dtype, zarr_format=zarr_format)
     shape_parsed = parse_shapelike(shape)
     chunk_key_encoding_parsed = _parse_chunk_key_encoding(
         chunk_key_encoding, zarr_format=zarr_format
@@ -3934,7 +3933,10 @@ async def init_array(
         await ensure_no_existing_node(store_path, zarr_format=zarr_format)
 
     shard_shape_parsed, chunk_shape_parsed = _auto_partition(
-        array_shape=shape_parsed, shard_shape=shards, chunk_shape=chunks, dtype=dtype_parsed
+        array_shape=shape_parsed,
+        shard_shape=shards,
+        chunk_shape=chunks,
+        item_size=dtype_wrapped.unwrap().itemsize,
     )
     chunks_out: tuple[int, ...]
     meta: ArrayV2Metadata | ArrayV3Metadata
@@ -3950,9 +3952,8 @@ async def init_array(
             raise ValueError("Zarr format 2 arrays do not support `serializer`.")
 
         filters_parsed, compressor_parsed = _parse_chunk_encoding_v2(
-            compressor=compressors, filters=filters, dtype=np.dtype(dtype)
+            compressor=compressors, filters=filters, dtype=dtype_wrapped
         )
-
         if dimension_names is not None:
             raise ValueError("Zarr format 2 arrays do not support dimension names.")
         if order is None:
@@ -3962,7 +3963,7 @@ async def init_array(
 
         meta = AsyncArray._create_metadata_v2(
             shape=shape_parsed,
-            dtype=dtype_parsed,
+            dtype=dtype_wrapped,
             chunks=chunk_shape_parsed,
             dimension_separator=chunk_key_encoding_parsed.separator,
             fill_value=fill_value,
@@ -3976,7 +3977,7 @@ async def init_array(
             compressors=compressors,
             filters=filters,
             serializer=serializer,
-            dtype=dtype_parsed,
+            dtype=dtype_wrapped,
         )
         sub_codecs = cast(tuple[Codec, ...], (*array_array, array_bytes, *bytes_bytes))
         codecs_out: tuple[Codec, ...]
@@ -3991,7 +3992,7 @@ async def init_array(
             )
             sharding_codec.validate(
                 shape=chunk_shape_parsed,
-                dtype=dtype_parsed,
+                dtype=dtype_wrapped,
                 chunk_grid=RegularChunkGrid(chunk_shape=shard_shape_parsed),
             )
             codecs_out = (sharding_codec,)
@@ -4002,7 +4003,7 @@ async def init_array(
 
         meta = AsyncArray._create_metadata_v3(
             shape=shape_parsed,
-            dtype=dtype_parsed,
+            dtype=dtype_wrapped,
             fill_value=fill_value,
             chunk_shape=chunks_out,
             chunk_key_encoding=chunk_key_encoding_parsed,
@@ -4210,12 +4211,11 @@ def _parse_chunk_key_encoding(
 
 
 def _get_default_chunk_encoding_v3(
-    np_dtype: np.dtype[Any],
+    dtype: DTypeWrapper[Any, Any],
 ) -> tuple[tuple[ArrayArrayCodec, ...], ArrayBytesCodec, tuple[BytesBytesCodec, ...]]:
     """
     Get the default ArrayArrayCodecs, ArrayBytesCodec, and BytesBytesCodec for a given dtype.
     """
-    dtype = get_data_type_from_numpy(np_dtype)
 
     default_filters = zarr_config.get("array.v3_default_filters").get(dtype.kind)
     default_serializer = zarr_config.get("array.v3_default_serializer").get(dtype.kind)
@@ -4229,14 +4229,14 @@ def _get_default_chunk_encoding_v3(
 
 
 def _get_default_chunk_encoding_v2(
-    np_dtype: np.dtype[Any],
+    dtype: DTypeWrapper[Any, Any],
 ) -> tuple[tuple[numcodecs.abc.Codec, ...] | None, numcodecs.abc.Codec | None]:
     """
     Get the default chunk encoding for Zarr format 2 arrays, given a dtype
     """
 
-    compressor_dict = _default_compressor(np_dtype)
-    filter_dicts = _default_filters(np_dtype)
+    compressor_dict = _default_compressor(dtype)
+    filter_dicts = _default_filters(dtype)
 
     compressor = None
     if compressor_dict is not None:
@@ -4253,13 +4253,12 @@ def _parse_chunk_encoding_v2(
     *,
     compressor: CompressorsLike,
     filters: FiltersLike,
-    dtype: np.dtype[Any],
+    dtype: DTypeWrapper[Any, Any],
 ) -> tuple[tuple[numcodecs.abc.Codec, ...] | None, numcodecs.abc.Codec | None]:
     """
     Generate chunk encoding classes for Zarr format 2 arrays with optional defaults.
     """
     default_filters, default_compressor = _get_default_chunk_encoding_v2(dtype)
-
     _filters: tuple[numcodecs.abc.Codec, ...] | None
     _compressor: numcodecs.abc.Codec | None
 

src/zarr/core/array_spec.py

@@ -3,8 +3,6 @@
 from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Literal, Self, TypedDict, cast
 
-import numpy as np
-
 from zarr.core.common import (
     MemoryOrder,
     parse_bool,
@@ -13,10 +11,14 @@
     parse_shapelike,
 )
 from zarr.core.config import config as zarr_config
+from zarr.core.metadata.dtype import DTypeWrapper
+from zarr.registry import get_data_type_from_numpy
 
 if TYPE_CHECKING:
     from typing import NotRequired
 
+    import numpy.typing as npt
+
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import ChunkCoords
 
@@ -90,21 +92,25 @@ def parse_array_config(data: ArrayConfigLike | None) -> ArrayConfig:
 @dataclass(frozen=True)
 class ArraySpec:
     shape: ChunkCoords
-    dtype: np.dtype[Any]
+    dtype: DTypeWrapper[Any, Any]
     fill_value: Any
     config: ArrayConfig
     prototype: BufferPrototype
 
     def __init__(
         self,
         shape: ChunkCoords,
-        dtype: np.dtype[Any],
+        dtype: npt.DtypeLike | DTypeWrapper[Any, Any],
         fill_value: Any,
         config: ArrayConfig,
         prototype: BufferPrototype,
     ) -> None:
         shape_parsed = parse_shapelike(shape)
-        dtype_parsed = np.dtype(dtype)
+        if not isinstance(dtype, DTypeWrapper):
+            dtype_parsed = get_data_type_from_numpy(dtype)
+        else:
+            dtype_parsed = dtype
+
         fill_value_parsed = parse_fill_value(fill_value)
 
         object.__setattr__(self, "shape", shape_parsed)

src/zarr/core/buffer/cpu.py

@@ -10,6 +10,7 @@
 import numpy.typing as npt
 
 from zarr.core.buffer import core
+from zarr.core.metadata.dtype import DTypeWrapper
 from zarr.registry import (
     register_buffer,
     register_ndbuffer,
@@ -150,14 +151,18 @@ def create(
         cls,
         *,
         shape: Iterable[int],
-        dtype: npt.DTypeLike,
+        dtype: DTypeWrapper[Any, Any],
         order: Literal["C", "F"] = "C",
         fill_value: Any | None = None,
     ) -> Self:
         if fill_value is None:
             return cls(np.zeros(shape=tuple(shape), dtype=dtype, order=order))
         else:
-            return cls(np.full(shape=tuple(shape), fill_value=fill_value, dtype=dtype, order=order))
+            return cls(
+                np.full(
+                    shape=tuple(shape), fill_value=fill_value, dtype=dtype.unwrap(), order=order
+                )
+            )
 
     @classmethod
     def from_numpy_array(cls, array_like: npt.ArrayLike) -> Self: