diff --git a/pyproject.toml b/pyproject.toml index 93888a205c..966108eabc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -180,10 +180,6 @@ check_untyped_defs = false module = [ "zarr.v2.*", "zarr.abc.codec", - "zarr.codecs.bytes", - "zarr.codecs.pipeline", - "zarr.codecs.sharding", - "zarr.codecs.transpose", "zarr.array_v2", "zarr.array", "zarr.sync", diff --git a/src/zarr/codecs/bytes.py b/src/zarr/codecs/bytes.py index aa24c3167e..aa8634645d 100644 --- a/src/zarr/codecs/bytes.py +++ b/src/zarr/codecs/bytes.py @@ -3,9 +3,10 @@ from enum import Enum import sys -from typing import TYPE_CHECKING, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Union import numpy as np +import numpy.typing as npt from zarr.abc.codec import ArrayBytesCodec from zarr.codecs.registry import register_codec @@ -60,7 +61,7 @@ def evolve(self, array_spec: ArraySpec) -> Self: ) return self - def _get_byteorder(self, array: np.ndarray) -> Endian: + def _get_byteorder(self, array: npt.NDArray[Any]) -> Endian: if array.dtype.byteorder == "<": return Endian.little elif array.dtype.byteorder == ">": @@ -73,7 +74,7 @@ async def decode( chunk_bytes: BytesLike, chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, - ) -> np.ndarray: + ) -> npt.NDArray[Any]: if chunk_spec.dtype.itemsize > 0: if self.endian == Endian.little: prefix = "<" @@ -93,7 +94,7 @@ async def decode( async def encode( self, - chunk_array: np.ndarray, + chunk_array: npt.NDArray[Any], _chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, ) -> Optional[BytesLike]: diff --git a/src/zarr/codecs/pipeline.py b/src/zarr/codecs/pipeline.py index 4908ee8057..b1cdb6353a 100644 --- a/src/zarr/codecs/pipeline.py +++ b/src/zarr/codecs/pipeline.py @@ -1,7 +1,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Iterable -import numpy as np +from typing import TYPE_CHECKING, Any, Iterable +import numpy.typing as npt from dataclasses import dataclass from warnings import warn @@ -152,7 +152,7 @@ async def decode( chunk_bytes: BytesLike, array_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, - ) -> np.ndarray: + ) -> npt.NDArray[Any]: ( aa_codecs_with_spec, ab_codec_with_spec, @@ -176,7 +176,7 @@ async def decode_partial( selection: SliceSelection, chunk_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, - ) -> Optional[np.ndarray]: + ) -> Optional[npt.NDArray[Any]]: assert self.supports_partial_decode assert isinstance(self.array_bytes_codec, ArrayBytesCodecPartialDecodeMixin) return await self.array_bytes_codec.decode_partial( @@ -185,7 +185,7 @@ async def decode_partial( async def encode( self, - chunk_array: np.ndarray, + chunk_array: npt.NDArray[Any], array_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, ) -> Optional[BytesLike]: @@ -222,7 +222,7 @@ async def encode( async def encode_partial( self, store_path: StorePath, - chunk_array: np.ndarray, + chunk_array: npt.NDArray[Any], selection: SliceSelection, chunk_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index d4f8b7dfc9..090393ffef 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -1,11 +1,12 @@ from __future__ import annotations from enum import Enum -from typing import TYPE_CHECKING, Iterable, Mapping, NamedTuple, Union +from typing import TYPE_CHECKING, Any, Iterable, Mapping, NamedTuple, Union, Optional from dataclasses import dataclass, replace from functools import lru_cache import numpy as np +import numpy.typing as npt from zarr.abc.codec import ( Codec, ArrayBytesCodec, @@ -18,7 +19,9 @@ from zarr.codecs.registry import register_codec from zarr.common import ( ArraySpec, + BytesLike, ChunkCoordsLike, + ChunkCoords, concurrent_map, parse_enum, parse_named_configuration, @@ -39,14 +42,12 @@ ) if TYPE_CHECKING: - from typing import Awaitable, Callable, Dict, Iterator, List, Optional, Set, Tuple + from typing import Awaitable, Callable, Dict, Iterator, List, Set, Tuple from typing_extensions import Self from zarr.store import StorePath from zarr.common import ( JSON, - ChunkCoords, - BytesLike, SliceSelection, ) from zarr.config import RuntimeConfiguration @@ -65,7 +66,7 @@ def parse_index_location(data: JSON) -> ShardingCodecIndexLocation: class _ShardIndex(NamedTuple): # dtype uint64, shape (chunks_per_shard_0, chunks_per_shard_1, ..., 2) - offsets_and_lengths: np.ndarray + offsets_and_lengths: npt.NDArray[np.uint64] @property def chunks_per_shard(self) -> ChunkCoords: @@ -126,7 +127,10 @@ def create_empty(cls, chunks_per_shard: ChunkCoords) -> _ShardIndex: return cls(offsets_and_lengths) -class _ShardProxy(Mapping): +_ShardMapping = Mapping[ChunkCoords, Optional[BytesLike]] + + +class _ShardProxy(_ShardMapping): index: _ShardIndex buf: BytesLike @@ -175,7 +179,7 @@ def merge_with_morton_order( cls, chunks_per_shard: ChunkCoords, tombstones: Set[ChunkCoords], - *shard_dicts: Mapping[ChunkCoords, BytesLike], + *shard_dicts: _ShardMapping, ) -> _ShardBuilder: obj = cls.create_empty(chunks_per_shard) for chunk_coords in morton_order_iter(chunks_per_shard): @@ -303,7 +307,7 @@ async def decode( shard_bytes: BytesLike, shard_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, - ) -> np.ndarray: + ) -> npt.NDArray[Any]: # print("decode") shard_shape = shard_spec.shape chunk_shape = self.chunk_shape @@ -353,7 +357,7 @@ async def decode_partial( selection: SliceSelection, shard_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, - ) -> Optional[np.ndarray]: + ) -> Optional[npt.NDArray[Any]]: shard_shape = shard_spec.shape chunk_shape = self.chunk_shape chunks_per_shard = self._get_chunks_per_shard(shard_spec) @@ -375,7 +379,7 @@ async def decode_partial( all_chunk_coords = set(chunk_coords for chunk_coords, _, _ in indexed_chunks) # reading bytes of all requested chunks - shard_dict: Mapping[ChunkCoords, BytesLike] = {} + shard_dict: _ShardMapping = {} if self._is_total_shard(all_chunk_coords, chunks_per_shard): # read entire shard shard_dict_maybe = await self._load_full_shard_maybe(store_path, chunks_per_shard) @@ -423,7 +427,7 @@ async def _read_chunk( out_selection: SliceSelection, shard_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, - out: np.ndarray, + out: npt.NDArray[Any], ) -> None: chunk_spec = self._get_chunk_spec(shard_spec) chunk_bytes = shard_dict.get(chunk_coords, None) @@ -436,7 +440,7 @@ async def _read_chunk( async def encode( self, - shard_array: np.ndarray, + shard_array: npt.NDArray[Any], shard_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, ) -> Optional[BytesLike]: @@ -453,7 +457,7 @@ async def encode( ) async def _write_chunk( - shard_array: np.ndarray, + shard_array: npt.NDArray[Any], chunk_coords: ChunkCoords, chunk_selection: SliceSelection, out_selection: SliceSelection, @@ -498,7 +502,7 @@ async def _write_chunk( async def encode_partial( self, store_path: StorePath, - shard_array: np.ndarray, + shard_array: npt.NDArray[Any], selection: SliceSelection, shard_spec: ArraySpec, runtime_configuration: RuntimeConfiguration, diff --git a/src/zarr/codecs/transpose.py b/src/zarr/codecs/transpose.py index c63327f6fc..10f5d66743 100644 --- a/src/zarr/codecs/transpose.py +++ b/src/zarr/codecs/transpose.py @@ -1,5 +1,5 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Dict, Iterable, Union, cast +from typing import TYPE_CHECKING, Any, Dict, Iterable, Union, cast from dataclasses import dataclass, replace @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Optional, Tuple from typing_extensions import Self -import numpy as np +import numpy.typing as npt from zarr.abc.codec import ArrayArrayCodec from zarr.codecs.registry import register_codec @@ -75,10 +75,10 @@ def resolve_metadata(self, chunk_spec: ArraySpec) -> ArraySpec: async def decode( self, - chunk_array: np.ndarray, + chunk_array: npt.NDArray[Any], chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, - ) -> np.ndarray: + ) -> npt.NDArray[Any]: inverse_order = [0] * chunk_spec.ndim for x, i in enumerate(self.order): inverse_order[x] = i @@ -87,10 +87,10 @@ async def decode( async def encode( self, - chunk_array: np.ndarray, + chunk_array: npt.NDArray[Any], chunk_spec: ArraySpec, _runtime_configuration: RuntimeConfiguration, - ) -> Optional[np.ndarray]: + ) -> Optional[npt.NDArray[Any]]: chunk_array = chunk_array.transpose(self.order) return chunk_array