adjust type annotations

Author: d-v-b

src/zarr/api/asynchronous.py

@@ -990,7 +990,7 @@ async def create(
             chunks = shape
         default_filters, default_compressor = _get_default_chunk_encoding_v2(dtype_wrapped)
         if filters is None:
-            filters = default_filters
+            filters = default_filters  # type: ignore[assignment]
         if compressor is None:
             compressor = default_compressor
     elif zarr_format == 3 and chunk_shape is None:  # type: ignore[redundant-expr]

src/zarr/codecs/_v2.py

@@ -46,7 +46,7 @@ async def _decode_single(
         chunk = ensure_ndarray_like(chunk)
         # special case object dtype, because incorrect handling can lead to
         # segfaults and other bad things happening
-        if chunk_spec.dtype != object:
+        if chunk_spec.dtype.dtype_cls is not np.dtypes.ObjectDType:
             try:
                 chunk = chunk.view(chunk_spec.dtype.to_dtype())
             except TypeError:

src/zarr/core/_info.py

@@ -2,14 +2,14 @@
 
 import dataclasses
 import textwrap
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 if TYPE_CHECKING:
     import numcodecs.abc
 
     from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec
     from zarr.core.common import ZarrFormat
-    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
+    from zarr.core.dtype.wrapper import ZDType
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -80,7 +80,7 @@ class ArrayInfo:
 
     _type: Literal["Array"] = "Array"
     _zarr_format: ZarrFormat
-    _data_type: ZDType[_BaseDType, _BaseScalar]
+    _data_type: ZDType[Any, Any]
     _shape: tuple[int, ...]
     _shard_shape: tuple[int, ...] | None = None
     _chunk_shape: tuple[int, ...] | None = None

src/zarr/core/array.py

@@ -680,6 +680,8 @@ def _create_metadata_v3(
         """
         Create an instance of ArrayV3Metadata.
         """
+        filters: tuple[ArrayArrayCodec, ...]
+        compressors: tuple[BytesBytesCodec, ...]
 
         shape = parse_shapelike(shape)
         if codecs is None:
@@ -707,7 +709,7 @@ def _create_metadata_v3(
             chunk_grid=chunk_grid_parsed,
             chunk_key_encoding=chunk_key_encoding_parsed,
             fill_value=fill_value_parsed,
-            codecs=codecs_parsed,
+            codecs=codecs_parsed,  # type: ignore[arg-type]
             dimension_names=tuple(dimension_names) if dimension_names else None,
             attributes=attributes or {},
         )
@@ -1712,7 +1714,7 @@ def _info(
     ) -> Any:
         return ArrayInfo(
             _zarr_format=self.metadata.zarr_format,
-            _data_type=self.dtype,
+            _data_type=self._zdtype,
             _shape=self.shape,
             _order=self.order,
             _shard_shape=self.shards,

src/zarr/core/array_spec.py

@@ -17,7 +17,7 @@
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import ChunkCoords
-    from zarr.core.dtype.wrapper import ZDType, _BaseDType, _BaseScalar
+    from zarr.core.dtype.wrapper import ZDType
 
 
 class ArrayConfigParams(TypedDict):
@@ -89,15 +89,15 @@ def parse_array_config(data: ArrayConfigLike | None) -> ArrayConfig:
 @dataclass(frozen=True)
 class ArraySpec:
     shape: ChunkCoords
-    dtype: ZDType[_BaseDType, _BaseScalar]
+    dtype: ZDType[Any, Any]
     fill_value: Any
     config: ArrayConfig
     prototype: BufferPrototype
 
     def __init__(
         self,
         shape: ChunkCoords,
-        dtype: ZDType[_BaseDType, _BaseScalar],
+        dtype: ZDType[Any, Any],
         fill_value: Any,
         config: ArrayConfig,
         prototype: BufferPrototype,

src/zarr/core/dtype/__init__.py

@@ -84,7 +84,11 @@
     data_type_registry.register(dtype._zarr_v3_name, dtype)
 
 
+# TODO: find a better name for this function
 def get_data_type_from_native_dtype(dtype: npt.DTypeLike) -> ZDType[_BaseDType, _BaseScalar]:
+    """
+    Get a data type wrapper (an instance of ``ZDType``) from a native data type, e.g. a numpy dtype.
+    """
     data_type_registry.lazy_load()
     if not isinstance(dtype, np.dtype):
         if dtype in (str, "str"):

src/zarr/core/dtype/wrapper.py

@@ -22,7 +22,7 @@
 TDType = TypeVar("TDType", bound=_BaseDType)
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(frozen=True, kw_only=True, slots=True)
 class ZDType(Generic[TDType, TScalar], ABC):
     """
     Abstract base class for wrapping native array data types, e.g. numpy dtypes
@@ -62,7 +62,7 @@ def check_dtype(cls: type[Self], dtype: _BaseDType) -> TypeGuard[TDType]:
         return type(dtype) is cls.dtype_cls
 
     @classmethod
-    def from_dtype(cls: type[Self], dtype: TDType) -> Self:
+    def from_dtype(cls: type[Self], dtype: _BaseDType) -> Self:
         """
         Wrap a dtype object.
 

src/zarr/core/metadata/v3.py

@@ -109,7 +109,8 @@ def validate_codecs(codecs: tuple[Codec, ...], dtype: ZDType[_BaseDType, _BaseSc
     # we need to have special codecs if we are decoding vlen strings or bytestrings
     # TODO: use codec ID instead of class name
     codec_class_name = abc.__class__.__name__
-    if isinstance(dtype, VariableLengthString) and not codec_class_name == "VLenUTF8Codec":
+    # TODO: Fix typing here
+    if isinstance(dtype, VariableLengthString) and not codec_class_name == "VLenUTF8Codec":  # type: ignore[unreachable]
         raise ValueError(
             f"For string dtype, ArrayBytesCodec must be `VLenUTF8Codec`, got `{codec_class_name}`."
         )
@@ -407,11 +408,11 @@ def to_dict(self) -> dict[str, JSON]:
 
         # TODO: replace the `to_dict` / `from_dict` on the `Metadata`` class with
         # to_json, from_json, and have ZDType inherit from `Metadata`
-        # until then, we have this hack here
+        # until then, we have this hack here, which relies on the fact that to_dict will pass through
+        # any non-`Metadata` fields as-is.
         dtype_meta = out_dict["data_type"]
-
         if isinstance(dtype_meta, ZDType):
-            out_dict["data_type"] = dtype_meta.to_json(zarr_format=3)
+            out_dict["data_type"] = dtype_meta.to_json(zarr_format=3)  # type: ignore[unreachable]
 
         return out_dict
 

tests/test_array.py

@@ -39,6 +39,7 @@
 from zarr.core.chunk_grids import _auto_partition
 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.group import AsyncGroup
 from zarr.core.indexing import BasicIndexer, ceildiv
 from zarr.core.sync import sync
@@ -448,7 +449,7 @@ def test_info_v2(self, chunks: tuple[int, int], shards: tuple[int, int] | None)
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=2,
-            _data_type=arr.dtype,
+            _data_type=arr._async_array._zdtype,
             _shape=(8, 8),
             _chunk_shape=chunks,
             _shard_shape=None,
@@ -465,7 +466,7 @@ def test_info_v3(self, chunks: tuple[int, int], shards: tuple[int, int] | None)
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=3,
-            _data_type=arr.dtype,
+            _data_type=arr._async_array._zdtype,
             _shape=(8, 8),
             _chunk_shape=chunks,
             _shard_shape=shards,
@@ -490,7 +491,7 @@ def test_info_complete(self, chunks: tuple[int, int], shards: tuple[int, int] |
         result = arr.info_complete()
         expected = ArrayInfo(
             _zarr_format=3,
-            _data_type=arr.dtype,
+            _data_type=arr._async_array._zdtype,
             _shape=(8, 8),
             _chunk_shape=chunks,
             _shard_shape=shards,
@@ -525,7 +526,7 @@ async def test_info_v2_async(
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=2,
-            _data_type=np.dtype("float64"),
+            _data_type=Float64(),
             _shape=(8, 8),
             _chunk_shape=(2, 2),
             _shard_shape=None,
@@ -550,7 +551,7 @@ async def test_info_v3_async(
         result = arr.info
         expected = ArrayInfo(
             _zarr_format=3,
-            _data_type=arr.dtype,
+            _data_type=arr._zdtype,
             _shape=(8, 8),
             _chunk_shape=chunks,
             _shard_shape=shards,
@@ -577,7 +578,7 @@ async def test_info_complete_async(
         result = await arr.info_complete()
         expected = ArrayInfo(
             _zarr_format=3,
-            _data_type=arr.dtype,
+            _data_type=arr._zdtype,
             _shape=(8, 8),
             _chunk_shape=chunks,
             _shard_shape=shards,

tests/test_info.py

@@ -1,11 +1,11 @@
 import textwrap
 
-import numpy as np
 import pytest
 
 from zarr.codecs.bytes import BytesCodec
 from zarr.core._info import ArrayInfo, GroupInfo, human_readable_size
 from zarr.core.common import ZarrFormat
+from zarr.core.dtype._numpy import Int32
 
 ZARR_FORMATS = [2, 3]
 
@@ -53,7 +53,7 @@ def test_group_info_complete(zarr_format: ZarrFormat) -> None:
 def test_array_info(zarr_format: ZarrFormat) -> None:
     info = ArrayInfo(
         _zarr_format=zarr_format,
-        _data_type=np.dtype("int32"),
+        _data_type=Int32(),
         _shape=(100, 100),
         _chunk_shape=(10, 100),
         _order="C",
@@ -91,7 +91,7 @@ def test_array_info_complete(
     ) = bytes_things
     info = ArrayInfo(
         _zarr_format=zarr_format,
-        _data_type=np.dtype("int32"),
+        _data_type=Int32(),
         _shape=(100, 100),
         _chunk_shape=(10, 100),
         _order="C",