refactor dtypewrapper -> zdtype

Author: d-v-b

src/zarr/abc/codec.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
 
 from zarr.abc.metadata import Metadata
 from zarr.core.buffer import Buffer, NDBuffer
@@ -12,11 +12,10 @@
     from collections.abc import Awaitable, Callable, Iterable
     from typing import Self
 
-    import numpy as np
-
     from zarr.abc.store import ByteGetter, ByteSetter
     from zarr.core.array_spec import ArraySpec
     from zarr.core.chunk_grids import ChunkGrid
+    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
     from zarr.core.indexing import SelectorTuple
 
 __all__ = [
@@ -93,7 +92,13 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         """
         return self
 
-    def validate(self, *, shape: ChunkCoords, dtype: np.dtype[Any], chunk_grid: ChunkGrid) -> None:
+    def validate(
+        self,
+        *,
+        shape: ChunkCoords,
+        dtype: ZDType[_BaseDType, _BaseScalar],
+        chunk_grid: ChunkGrid,
+    ) -> None:
         """Validates that the codec configuration is compatible with the array metadata.
         Raises errors when the codec configuration is not compatible.
 
@@ -285,7 +290,9 @@ def supports_partial_decode(self) -> bool: ...
     def supports_partial_encode(self) -> bool: ...
 
     @abstractmethod
-    def validate(self, *, shape: ChunkCoords, dtype: np.dtype[Any], chunk_grid: ChunkGrid) -> None:
+    def validate(
+        self, *, shape: ChunkCoords, dtype: ZDType[_BaseDType, _BaseScalar], chunk_grid: ChunkGrid
+    ) -> None:
         """Validates that all codec configurations are compatible with the array metadata.
         Raises errors when a codec configuration is not compatible.
 

src/zarr/api/asynchronous.py

@@ -28,7 +28,7 @@
     _warn_order_kwarg,
     _warn_write_empty_chunks_kwarg,
 )
-from zarr.core.dtype import get_data_type_from_numpy
+from zarr.core.dtype import get_data_type_from_native_dtype
 from zarr.core.group import (
     AsyncGroup,
     ConsolidatedMetadata,
@@ -433,7 +433,7 @@ 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)
+    zarr_dtype = get_data_type_from_native_dtype(arr.dtype)
     new = await AsyncArray._create(
         store_path,
         zarr_format=zarr_format,
@@ -984,7 +984,7 @@ async def create(
         _handle_zarr_version_or_format(zarr_version=zarr_version, zarr_format=zarr_format)
         or _default_zarr_format()
     )
-    dtype_wrapped = get_data_type_from_numpy(dtype)
+    dtype_wrapped = get_data_type_from_native_dtype(dtype)
     if zarr_format == 2:
         if chunks is None:
             chunks = shape

src/zarr/codecs/bytes.py

@@ -3,20 +3,21 @@
 import sys
 from dataclasses import dataclass, replace
 from enum import Enum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
 from zarr.abc.codec import ArrayBytesCodec
 from zarr.core.buffer import Buffer, NDArrayLike, NDBuffer
 from zarr.core.common import JSON, parse_enum, parse_named_configuration
-from zarr.core.dtype.common import endianness_to_numpy_str
+from zarr.core.dtype._numpy import endianness_to_numpy_str
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
     from typing import Self
 
     from zarr.core.array_spec import ArraySpec
+    from zarr.core.dtype.common import Endianness
 
 
 class Endian(Enum):
@@ -73,7 +74,9 @@ async def _decode_single(
     ) -> NDBuffer:
         assert isinstance(chunk_bytes, Buffer)
         # TODO: remove endianness enum in favor of literal union
-        endian_str = self.endian.value if self.endian is not None else None
+        endian_str = cast(
+            "Endianness | None", self.endian.value if self.endian is not None else None
+        )
         dtype = chunk_spec.dtype.to_dtype().newbyteorder(endianness_to_numpy_str(endian_str))
 
         as_array_like = chunk_bytes.as_array_like()

src/zarr/codecs/sharding.py

@@ -43,6 +43,7 @@
     parse_shapelike,
     product,
 )
+from zarr.core.dtype._numpy import UInt64
 from zarr.core.indexing import (
     BasicIndexer,
     SelectorTuple,
@@ -58,7 +59,7 @@
     from typing import Self
 
     from zarr.core.common import JSON
-    from zarr.core.dtype.wrapper import DTypeWrapper
+    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
 
 MAX_UINT_64 = 2**64 - 1
 ShardMapping = Mapping[ChunkCoords, Buffer]
@@ -405,7 +406,11 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self:
         return self
 
     def validate(
-        self, *, shape: ChunkCoords, dtype: DTypeWrapper[Any, Any], chunk_grid: ChunkGrid
+        self,
+        *,
+        shape: ChunkCoords,
+        dtype: ZDType[_BaseDType, _BaseScalar],
+        chunk_grid: ChunkGrid,
     ) -> None:
         if len(self.chunk_shape) != len(shape):
             raise ValueError(
@@ -443,7 +448,10 @@ async def _decode_single(
 
         # setup output array
         out = chunk_spec.prototype.nd_buffer.create(
-            shape=shard_shape, dtype=shard_spec.dtype, order=shard_spec.order, fill_value=0
+            shape=shard_shape,
+            dtype=shard_spec.dtype.to_dtype(),
+            order=shard_spec.order,
+            fill_value=0,
         )
         shard_dict = await _ShardReader.from_bytes(shard_bytes, self, chunks_per_shard)
 
@@ -685,7 +693,7 @@ def _shard_index_size(self, chunks_per_shard: ChunkCoords) -> int:
     def _get_index_chunk_spec(self, chunks_per_shard: ChunkCoords) -> ArraySpec:
         return ArraySpec(
             shape=chunks_per_shard + (2,),
-            dtype=np.dtype("<u8"),
+            dtype=UInt64(endianness="little"),
             fill_value=MAX_UINT_64,
             config=ArrayConfig(
                 order="C", write_empty_chunks=False

src/zarr/codecs/transpose.py

@@ -12,10 +12,11 @@
 from zarr.registry import register_codec
 
 if TYPE_CHECKING:
-    from typing import Any, Self
+    from typing import Self
 
     from zarr.core.buffer import NDBuffer
     from zarr.core.chunk_grids import ChunkGrid
+    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
 
 
 def parse_transpose_order(data: JSON | Iterable[int]) -> tuple[int, ...]:
@@ -45,7 +46,12 @@ def from_dict(cls, data: dict[str, JSON]) -> Self:
     def to_dict(self) -> dict[str, JSON]:
         return {"name": "transpose", "configuration": {"order": tuple(self.order)}}
 
-    def validate(self, shape: tuple[int, ...], dtype: np.dtype[Any], chunk_grid: ChunkGrid) -> None:
+    def validate(
+        self,
+        shape: tuple[int, ...],
+        dtype: ZDType[_BaseDType, _BaseScalar],
+        chunk_grid: ChunkGrid,
+    ) -> None:
         if len(self.order) != len(shape):
             raise ValueError(
                 f"The `order` tuple needs have as many entries as there are dimensions in the array. Got {self.order}."

src/zarr/core/_info.py

@@ -1,13 +1,16 @@
+from __future__ import annotations
+
 import dataclasses
 import textwrap
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
-import numcodecs.abc
-import numpy as np
+if TYPE_CHECKING:
+    import numcodecs.abc
+    import numpy as np
 
-from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
-from zarr.core.common import ZarrFormat
-from zarr.core.dtype.wrapper import DTypeWrapper
+    from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
+    from zarr.core.common import ZarrFormat
+    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -78,7 +81,7 @@ class ArrayInfo:
 
     _type: Literal["Array"] = "Array"
     _zarr_format: ZarrFormat
-    _data_type: np.dtype[Any] | DTypeWrapper
+    _data_type: np.dtype[Any] | ZDType[_BaseDType, _BaseScalar]
     _shape: tuple[int, ...]
     _shard_shape: tuple[int, ...] | None = None
     _chunk_shape: tuple[int, ...] | None = None