diff --git a/changes/3374.misc.rst b/changes/3374.misc.rst new file mode 100644 index 0000000000..19cefd8d60 --- /dev/null +++ b/changes/3374.misc.rst @@ -0,0 +1 @@ +Replaces usage of the ``zarr.core.common.ChunkCoords`` typealias with ``tuple[int, ...]``. \ No newline at end of file diff --git a/src/zarr/abc/codec.py b/src/zarr/abc/codec.py index f8a5447a70..fa3650b560 100644 --- a/src/zarr/abc/codec.py +++ b/src/zarr/abc/codec.py @@ -5,7 +5,7 @@ from zarr.abc.metadata import Metadata from zarr.core.buffer import Buffer, NDBuffer -from zarr.core.common import ChunkCoords, concurrent_map +from zarr.core.common import concurrent_map from zarr.core.config import config if TYPE_CHECKING: @@ -96,7 +96,7 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: def validate( self, *, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGrid, ) -> None: @@ -105,7 +105,7 @@ def validate( Parameters ---------- - shape : ChunkCoords + shape : tuple[int, ...] The array shape dtype : np.dtype[Any] The array data type @@ -311,14 +311,18 @@ def supports_partial_encode(self) -> bool: ... @abstractmethod def validate( - self, *, shape: ChunkCoords, dtype: ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGrid + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGrid, ) -> None: """Validates that all codec configurations are compatible with the array metadata. Raises errors when a codec configuration is not compatible. Parameters ---------- - shape : ChunkCoords + shape : tuple[int, ...] The array shape dtype : np.dtype[Any] The array data type diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 78b68caf73..58f1f80202 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -24,7 +24,6 @@ from zarr.core.common import ( JSON, AccessModeLiteral, - ChunkCoords, DimensionNames, MemoryOrder, ZarrFormat, @@ -107,7 +106,7 @@ def _infer_overwrite(mode: AccessModeLiteral) -> bool: return mode in _OVERWRITE_MODES -def _get_shape_chunks(a: ArrayLike | Any) -> tuple[ChunkCoords | None, ChunkCoords | None]: +def _get_shape_chunks(a: ArrayLike | Any) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None]: """Helper function to get the shape and chunks from an array-like object""" shape = None chunks = None @@ -865,9 +864,9 @@ async def open_group( async def create( - shape: ChunkCoords | int, + shape: tuple[int, ...] | int, *, # Note: this is a change from v2 - chunks: ChunkCoords | int | bool | None = None, + chunks: tuple[int, ...] | int | bool | None = None, dtype: ZDTypeLike | None = None, compressor: CompressorLike = "auto", fill_value: Any | None = DEFAULT_FILL_VALUE, @@ -889,7 +888,7 @@ async def create( meta_array: Any | None = None, # TODO: need type attributes: dict[str, JSON] | None = None, # v3 only - chunk_shape: ChunkCoords | int | None = None, + chunk_shape: tuple[int, ...] | int | None = None, chunk_key_encoding: ( ChunkKeyEncoding | tuple[Literal["default"], Literal[".", "/"]] @@ -1074,7 +1073,7 @@ async def create( async def empty( - shape: ChunkCoords, **kwargs: Any + shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an empty array with the specified shape. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1126,7 +1125,7 @@ async def empty_like( # TODO: add type annotations for fill_value and kwargs async def full( - shape: ChunkCoords, fill_value: Any, **kwargs: Any + shape: tuple[int, ...], fill_value: Any, **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with `fill_value` being used as the default value for uninitialized portions of the array. @@ -1173,7 +1172,7 @@ async def full_like( async def ones( - shape: ChunkCoords, **kwargs: Any + shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with one being used as the default value for uninitialized portions of the array. @@ -1296,7 +1295,7 @@ async def open_like( async def zeros( - shape: ChunkCoords, **kwargs: Any + shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with zero being used as the default value for uninitialized portions of the array. diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index ed1ae2cf2a..d4049f003b 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -33,7 +33,6 @@ from zarr.core.common import ( JSON, AccessModeLiteral, - ChunkCoords, DimensionNames, MemoryOrder, ShapeLike, @@ -598,9 +597,9 @@ def create_group( # TODO: add type annotations for kwargs def create( - shape: ChunkCoords | int, + shape: tuple[int, ...] | int, *, # Note: this is a change from v2 - chunks: ChunkCoords | int | bool | None = None, + chunks: tuple[int, ...] | int | bool | None = None, dtype: ZDTypeLike | None = None, compressor: CompressorLike = "auto", fill_value: Any | None = DEFAULT_FILL_VALUE, # TODO: need type @@ -622,7 +621,7 @@ def create( meta_array: Any | None = None, # TODO: need type attributes: dict[str, JSON] | None = None, # v3 only - chunk_shape: ChunkCoords | int | None = None, + chunk_shape: tuple[int, ...] | int | None = None, chunk_key_encoding: ( ChunkKeyEncoding | tuple[Literal["default"], Literal[".", "/"]] @@ -755,7 +754,7 @@ def create_array( shape: ShapeLike | None = None, dtype: ZDTypeLike | None = None, data: np.ndarray[Any, np.dtype[Any]] | None = None, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -782,7 +781,7 @@ def create_array( name : str or None, optional The name of the array within the store. If ``name`` is ``None``, the array will be located at the root of the store. - shape : ChunkCoords, optional + shape : tuple[int, ...], optional Shape of the array. Can be ``None`` if ``data`` is provided. dtype : ZDTypeLike, optional Data type of the array. Can be ``None`` if ``data`` is provided. @@ -790,10 +789,10 @@ def create_array( Array-like data to use for initializing the array. If this parameter is provided, the ``shape`` and ``dtype`` parameters must be identical to ``data.shape`` and ``data.dtype``, or ``None``. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that @@ -921,7 +920,7 @@ def from_array( data: Array | npt.ArrayLike, write_data: bool = True, name: str | None = None, - chunks: Literal["auto", "keep"] | ChunkCoords = "keep", + chunks: Literal["auto", "keep"] | tuple[int, ...] = "keep", shards: ShardsLike | None | Literal["keep"] = "keep", filters: FiltersLike | Literal["keep"] = "keep", compressors: CompressorsLike | Literal["keep"] = "keep", @@ -951,22 +950,22 @@ def from_array( name : str or None, optional The name of the array within the store. If ``name`` is ``None``, the array will be located at the root of the store. - chunks : ChunkCoords or "auto" or "keep", optional + chunks : tuple[int, ...] or "auto" or "keep", optional Chunk shape of the array. Following values are supported: - "auto": Automatically determine the chunk shape based on the array's shape and dtype. - "keep": Retain the chunk shape of the data array if it is a zarr Array. - - ChunkCoords: A tuple of integers representing the chunk shape. + - tuple[int, ...]: A tuple of integers representing the chunk shape. If not specified, defaults to "keep" if data is a zarr Array, otherwise "auto". - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. Following values are supported: - "auto": Automatically determine the shard shape based on the array's shape and chunk shape. - "keep": Retain the shard shape of the data array if it is a zarr Array. - - ChunkCoords: A tuple of integers representing the shard shape. + - tuple[int, ...]: A tuple of integers representing the shard shape. - None: No sharding. If not specified, defaults to "keep" if data is a zarr Array, otherwise None. @@ -1129,7 +1128,7 @@ def from_array( # TODO: add type annotations for kwargs -def empty(shape: ChunkCoords, **kwargs: Any) -> Array: +def empty(shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an empty array with the specified shape. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1182,7 +1181,7 @@ def empty_like(a: ArrayLike, **kwargs: Any) -> Array: # TODO: add type annotations for kwargs and fill_value -def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> Array: +def full(shape: tuple[int, ...], fill_value: Any, **kwargs: Any) -> Array: """Create an array with a default fill value. Parameters @@ -1223,7 +1222,7 @@ def full_like(a: ArrayLike, **kwargs: Any) -> Array: # TODO: add type annotations for kwargs -def ones(shape: ChunkCoords, **kwargs: Any) -> Array: +def ones(shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an array with a fill value of one. Parameters @@ -1325,7 +1324,7 @@ def open_like(a: ArrayLike, path: str, **kwargs: Any) -> Array: # TODO: add type annotations for kwargs -def zeros(shape: ChunkCoords, **kwargs: Any) -> Array: +def zeros(shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an array with a fill value of zero. Parameters diff --git a/src/zarr/codecs/sharding.py b/src/zarr/codecs/sharding.py index 888d258649..2b9b2259d8 100644 --- a/src/zarr/codecs/sharding.py +++ b/src/zarr/codecs/sharding.py @@ -36,8 +36,7 @@ ) from zarr.core.chunk_grids import ChunkGrid, RegularChunkGrid from zarr.core.common import ( - ChunkCoords, - ChunkCoordsLike, + ShapeLike, parse_enum, parse_named_configuration, parse_shapelike, @@ -62,8 +61,8 @@ from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType MAX_UINT_64 = 2**64 - 1 -ShardMapping = Mapping[ChunkCoords, Buffer] -ShardMutableMapping = MutableMapping[ChunkCoords, Buffer] +ShardMapping = Mapping[tuple[int, ...], Buffer] +ShardMutableMapping = MutableMapping[tuple[int, ...], Buffer] class ShardingCodecIndexLocation(Enum): @@ -82,7 +81,7 @@ def parse_index_location(data: object) -> ShardingCodecIndexLocation: @dataclass(frozen=True) class _ShardingByteGetter(ByteGetter): shard_dict: ShardMapping - chunk_coords: ChunkCoords + chunk_coords: tuple[int, ...] async def get( self, prototype: BufferPrototype, byte_range: ByteRequest | None = None @@ -114,12 +113,12 @@ class _ShardIndex(NamedTuple): offsets_and_lengths: npt.NDArray[np.uint64] @property - def chunks_per_shard(self) -> ChunkCoords: + def chunks_per_shard(self) -> tuple[int, ...]: result = tuple(self.offsets_and_lengths.shape[0:-1]) # The cast is required until https://github.com/numpy/numpy/pull/27211 is merged - return cast("ChunkCoords", result) + return cast("tuple[int, ...]", result) - def _localize_chunk(self, chunk_coords: ChunkCoords) -> ChunkCoords: + def _localize_chunk(self, chunk_coords: tuple[int, ...]) -> tuple[int, ...]: return tuple( chunk_i % shard_i for chunk_i, shard_i in zip(chunk_coords, self.offsets_and_lengths.shape, strict=False) @@ -131,7 +130,7 @@ def is_all_empty(self) -> bool: def get_full_chunk_map(self) -> npt.NDArray[np.bool_]: return np.not_equal(self.offsets_and_lengths[..., 0], MAX_UINT_64) - def get_chunk_slice(self, chunk_coords: ChunkCoords) -> tuple[int, int] | None: + def get_chunk_slice(self, chunk_coords: tuple[int, ...]) -> tuple[int, int] | None: localized_chunk = self._localize_chunk(chunk_coords) chunk_start, chunk_len = self.offsets_and_lengths[localized_chunk] if (chunk_start, chunk_len) == (MAX_UINT_64, MAX_UINT_64): @@ -139,7 +138,7 @@ def get_chunk_slice(self, chunk_coords: ChunkCoords) -> tuple[int, int] | None: else: return (int(chunk_start), int(chunk_start + chunk_len)) - def set_chunk_slice(self, chunk_coords: ChunkCoords, chunk_slice: slice | None) -> None: + def set_chunk_slice(self, chunk_coords: tuple[int, ...], chunk_slice: slice | None) -> None: localized_chunk = self._localize_chunk(chunk_coords) if chunk_slice is None: self.offsets_and_lengths[localized_chunk] = (MAX_UINT_64, MAX_UINT_64) @@ -171,7 +170,7 @@ def is_dense(self, chunk_byte_length: int) -> bool: ) @classmethod - def create_empty(cls, chunks_per_shard: ChunkCoords) -> _ShardIndex: + def create_empty(cls, chunks_per_shard: tuple[int, ...]) -> _ShardIndex: offsets_and_lengths = np.zeros(chunks_per_shard + (2,), dtype=" _ShardReader: shard_index_size = codec._shard_index_size(chunks_per_shard) obj = cls() @@ -198,7 +197,7 @@ async def from_bytes( @classmethod def create_empty( - cls, chunks_per_shard: ChunkCoords, buffer_prototype: BufferPrototype | None = None + cls, chunks_per_shard: tuple[int, ...], buffer_prototype: BufferPrototype | None = None ) -> _ShardReader: if buffer_prototype is None: buffer_prototype = default_buffer_prototype() @@ -208,7 +207,7 @@ def create_empty( obj.index = index return obj - def __getitem__(self, chunk_coords: ChunkCoords) -> Buffer: + def __getitem__(self, chunk_coords: tuple[int, ...]) -> Buffer: chunk_byte_slice = self.index.get_chunk_slice(chunk_coords) if chunk_byte_slice: return self.buf[chunk_byte_slice[0] : chunk_byte_slice[1]] @@ -217,7 +216,7 @@ def __getitem__(self, chunk_coords: ChunkCoords) -> Buffer: def __len__(self) -> int: return int(self.index.offsets_and_lengths.size / 2) - def __iter__(self) -> Iterator[ChunkCoords]: + def __iter__(self) -> Iterator[tuple[int, ...]]: return c_order_iter(self.index.offsets_and_lengths.shape[:-1]) def is_empty(self) -> bool: @@ -231,8 +230,8 @@ class _ShardBuilder(_ShardReader, ShardMutableMapping): @classmethod def merge_with_morton_order( cls, - chunks_per_shard: ChunkCoords, - tombstones: set[ChunkCoords], + chunks_per_shard: tuple[int, ...], + tombstones: set[tuple[int, ...]], *shard_dicts: ShardMapping, ) -> _ShardBuilder: obj = cls.create_empty(chunks_per_shard) @@ -248,7 +247,7 @@ def merge_with_morton_order( @classmethod def create_empty( - cls, chunks_per_shard: ChunkCoords, buffer_prototype: BufferPrototype | None = None + cls, chunks_per_shard: tuple[int, ...], buffer_prototype: BufferPrototype | None = None ) -> _ShardBuilder: if buffer_prototype is None: buffer_prototype = default_buffer_prototype() @@ -257,13 +256,13 @@ def create_empty( obj.index = _ShardIndex.create_empty(chunks_per_shard) return obj - def __setitem__(self, chunk_coords: ChunkCoords, value: Buffer) -> None: + def __setitem__(self, chunk_coords: tuple[int, ...], value: Buffer) -> None: chunk_start = len(self.buf) chunk_length = len(value) self.buf += value self.index.set_chunk_slice(chunk_coords, slice(chunk_start, chunk_start + chunk_length)) - def __delitem__(self, chunk_coords: ChunkCoords) -> None: + def __delitem__(self, chunk_coords: tuple[int, ...]) -> None: raise NotImplementedError async def finalize( @@ -286,24 +285,24 @@ async def finalize( class _MergingShardBuilder(ShardMutableMapping): old_dict: _ShardReader new_dict: _ShardBuilder - tombstones: set[ChunkCoords] = field(default_factory=set) + tombstones: set[tuple[int, ...]] = field(default_factory=set) - def __getitem__(self, chunk_coords: ChunkCoords) -> Buffer: + def __getitem__(self, chunk_coords: tuple[int, ...]) -> Buffer: chunk_bytes_maybe = self.new_dict.get(chunk_coords) if chunk_bytes_maybe is not None: return chunk_bytes_maybe return self.old_dict[chunk_coords] - def __setitem__(self, chunk_coords: ChunkCoords, value: Buffer) -> None: + def __setitem__(self, chunk_coords: tuple[int, ...], value: Buffer) -> None: self.new_dict[chunk_coords] = value - def __delitem__(self, chunk_coords: ChunkCoords) -> None: + def __delitem__(self, chunk_coords: tuple[int, ...]) -> None: self.tombstones.add(chunk_coords) def __len__(self) -> int: return self.old_dict.__len__() - def __iter__(self) -> Iterator[ChunkCoords]: + def __iter__(self) -> Iterator[tuple[int, ...]]: return self.old_dict.__iter__() def is_empty(self) -> bool: @@ -335,7 +334,7 @@ class ShardingCodec( ): """Sharding codec""" - chunk_shape: ChunkCoords + chunk_shape: tuple[int, ...] codecs: tuple[Codec, ...] index_codecs: tuple[Codec, ...] index_location: ShardingCodecIndexLocation = ShardingCodecIndexLocation.end @@ -343,7 +342,7 @@ class ShardingCodec( def __init__( self, *, - chunk_shape: ChunkCoordsLike, + chunk_shape: ShapeLike, codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(),), index_codecs: Iterable[Codec | dict[str, JSON]] = (BytesCodec(), Crc32cCodec()), index_location: ShardingCodecIndexLocation | str = ShardingCodecIndexLocation.end, @@ -413,7 +412,7 @@ def evolve_from_array_spec(self, array_spec: ArraySpec) -> Self: def validate( self, *, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGrid, ) -> None: @@ -646,14 +645,14 @@ async def _encode_partial_single( ) def _is_total_shard( - self, all_chunk_coords: set[ChunkCoords], chunks_per_shard: ChunkCoords + self, all_chunk_coords: set[tuple[int, ...]], chunks_per_shard: tuple[int, ...] ) -> bool: return len(all_chunk_coords) == product(chunks_per_shard) and all( chunk_coords in all_chunk_coords for chunk_coords in c_order_iter(chunks_per_shard) ) async def _decode_shard_index( - self, index_bytes: Buffer, chunks_per_shard: ChunkCoords + self, index_bytes: Buffer, chunks_per_shard: tuple[int, ...] ) -> _ShardIndex: index_array = next( iter( @@ -686,7 +685,7 @@ async def _encode_shard_index(self, index: _ShardIndex) -> Buffer: assert isinstance(index_bytes, Buffer) return index_bytes - def _shard_index_size(self, chunks_per_shard: ChunkCoords) -> int: + def _shard_index_size(self, chunks_per_shard: tuple[int, ...]) -> int: return ( get_pipeline_class() .from_codecs(self.index_codecs) @@ -695,7 +694,7 @@ def _shard_index_size(self, chunks_per_shard: ChunkCoords) -> int: ) ) - def _get_index_chunk_spec(self, chunks_per_shard: ChunkCoords) -> ArraySpec: + def _get_index_chunk_spec(self, chunks_per_shard: tuple[int, ...]) -> ArraySpec: return ArraySpec( shape=chunks_per_shard + (2,), dtype=UInt64(endianness="little"), @@ -715,7 +714,7 @@ def _get_chunk_spec(self, shard_spec: ArraySpec) -> ArraySpec: prototype=shard_spec.prototype, ) - def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> ChunkCoords: + def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> tuple[int, ...]: return tuple( s // c for s, c in zip( @@ -726,7 +725,7 @@ def _get_chunks_per_shard(self, shard_spec: ArraySpec) -> ChunkCoords: ) async def _load_shard_index_maybe( - self, byte_getter: ByteGetter, chunks_per_shard: ChunkCoords + self, byte_getter: ByteGetter, chunks_per_shard: tuple[int, ...] ) -> _ShardIndex | None: shard_index_size = self._shard_index_size(chunks_per_shard) if self.index_location == ShardingCodecIndexLocation.start: @@ -743,14 +742,14 @@ async def _load_shard_index_maybe( return None async def _load_shard_index( - self, byte_getter: ByteGetter, chunks_per_shard: ChunkCoords + self, byte_getter: ByteGetter, chunks_per_shard: tuple[int, ...] ) -> _ShardIndex: return ( await self._load_shard_index_maybe(byte_getter, chunks_per_shard) ) or _ShardIndex.create_empty(chunks_per_shard) async def _load_full_shard_maybe( - self, byte_getter: ByteGetter, prototype: BufferPrototype, chunks_per_shard: ChunkCoords + self, byte_getter: ByteGetter, prototype: BufferPrototype, chunks_per_shard: tuple[int, ...] ) -> _ShardReader | None: shard_bytes = await byte_getter.get(prototype=prototype) diff --git a/src/zarr/codecs/transpose.py b/src/zarr/codecs/transpose.py index d6310d38a4..92e7da81e1 100644 --- a/src/zarr/codecs/transpose.py +++ b/src/zarr/codecs/transpose.py @@ -8,7 +8,7 @@ from zarr.abc.codec import ArrayArrayCodec from zarr.core.array_spec import ArraySpec -from zarr.core.common import JSON, ChunkCoordsLike, parse_named_configuration +from zarr.core.common import JSON, parse_named_configuration from zarr.registry import register_codec if TYPE_CHECKING: @@ -35,7 +35,7 @@ class TransposeCodec(ArrayArrayCodec): order: tuple[int, ...] - def __init__(self, *, order: ChunkCoordsLike) -> None: + def __init__(self, *, order: Iterable[int]) -> None: order_parsed = parse_transpose_order(order) object.__setattr__(self, "order", order_parsed) diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 311a0eb986..7226bf6767 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -54,7 +54,6 @@ ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON, - ChunkCoords, DimensionNames, MemoryOrder, ShapeLike, @@ -492,7 +491,7 @@ async def create( The fill value of the array (default is None). attributes : dict[str, JSON], optional The attributes of the array (default is None). - chunk_shape : ChunkCoords, optional + chunk_shape : tuple[int, ...], optional The shape of the array's chunks Zarr format 3 only. Zarr format 2 arrays should use `chunks` instead. If not specified, default are guessed based on the shape and dtype. @@ -711,7 +710,7 @@ async def _create( def _create_metadata_v3( shape: ShapeLike, dtype: ZDType[TBaseDType, TBaseScalar], - chunk_shape: ChunkCoords, + chunk_shape: tuple[int, ...], fill_value: Any | None = DEFAULT_FILL_VALUE, chunk_key_encoding: ChunkKeyEncodingLike | None = None, codecs: Iterable[Codec | dict[str, JSON]] | None = None, @@ -766,7 +765,7 @@ async def _create_v3( *, shape: ShapeLike, dtype: ZDType[TBaseDType, TBaseScalar], - chunk_shape: ChunkCoords, + chunk_shape: tuple[int, ...], config: ArrayConfig, fill_value: Any | None = DEFAULT_FILL_VALUE, chunk_key_encoding: ( @@ -812,9 +811,9 @@ async def _create_v3( @staticmethod def _create_metadata_v2( - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], - chunks: ChunkCoords, + chunks: tuple[int, ...], order: MemoryOrder, dimension_separator: Literal[".", "/"] | None = None, fill_value: Any | None = DEFAULT_FILL_VALUE, @@ -849,9 +848,9 @@ async def _create_v2( cls, store_path: StorePath, *, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], - chunks: ChunkCoords, + chunks: tuple[int, ...], order: MemoryOrder, config: ArrayConfig, dimension_separator: Literal[".", "/"] | None = None, @@ -981,7 +980,7 @@ def ndim(self) -> int: return len(self.metadata.shape) @property - def shape(self) -> ChunkCoords: + def shape(self) -> tuple[int, ...]: """Returns the shape of the Array. Returns @@ -992,7 +991,7 @@ def shape(self) -> ChunkCoords: return self.metadata.shape @property - def chunks(self) -> ChunkCoords: + def chunks(self) -> tuple[int, ...]: """Returns the chunk shape of the Array. If sharding is used the inner chunk shape is returned. @@ -1001,13 +1000,13 @@ def chunks(self) -> ChunkCoords: Returns ------- - ChunkCoords: + tuple[int, ...]: The chunk shape of the Array. """ return self.metadata.chunks @property - def shards(self) -> ChunkCoords | None: + def shards(self) -> tuple[int, ...] | None: """Returns the shard shape of the Array. Returns None if sharding is not used. @@ -1016,7 +1015,7 @@ def shards(self) -> ChunkCoords | None: Returns ------- - ChunkCoords: + tuple[int, ...]: The shard shape of the Array. """ return self.metadata.shards @@ -1185,7 +1184,7 @@ def basename(self) -> str: return self.name.split("/")[-1] @property - def cdata_shape(self) -> ChunkCoords: + def cdata_shape(self) -> tuple[int, ...]: """ The shape of the chunk grid for this array. @@ -1239,7 +1238,7 @@ async def nbytes_stored(self) -> int: def _iter_chunk_coords( self, *, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[ChunkCoords]: + ) -> Iterator[tuple[int, ...]]: """ Create an iterator over the coordinates of chunks in chunk grid space. If the `origin` keyword is used, iteration will start at the chunk index specified by `origin`. @@ -1257,7 +1256,7 @@ def _iter_chunk_coords( Yields ------ - chunk_coords: ChunkCoords + chunk_coords: tuple[int, ...] The coordinates of each chunk in the selection. """ return _iter_grid(self.cdata_shape, origin=origin, selection_shape=selection_shape) @@ -1627,7 +1626,7 @@ async def resize(self, new_shape: ShapeLike, delete_outside_chunks: bool = True) Parameters ---------- - new_shape : ChunkCoords + new_shape : tuple[int, ...] The desired new shape of the array. delete_outside_chunks : bool, optional @@ -1675,7 +1674,7 @@ async def _delete_key(key: str) -> None: # Update metadata (in place) object.__setattr__(self, "metadata", new_metadata) - async def append(self, data: npt.ArrayLike, axis: int = 0) -> ChunkCoords: + async def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Append `data` to `axis`. Parameters @@ -1861,13 +1860,13 @@ def create( store: StoreLike, *, # v2 and v3 - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDTypeLike, zarr_format: ZarrFormat = 3, fill_value: Any | None = DEFAULT_FILL_VALUE, attributes: dict[str, JSON] | None = None, # v3 only - chunk_shape: ChunkCoords | None = None, + chunk_shape: tuple[int, ...] | None = None, chunk_key_encoding: ( ChunkKeyEncoding | tuple[Literal["default"], Literal[".", "/"]] @@ -1877,7 +1876,7 @@ def create( codecs: Iterable[Codec | dict[str, JSON]] | None = None, dimension_names: DimensionNames = None, # v2 only - chunks: ChunkCoords | None = None, + chunks: tuple[int, ...] | None = None, dimension_separator: Literal[".", "/"] | None = None, order: MemoryOrder | None = None, filters: list[dict[str, JSON]] | None = None, @@ -1895,11 +1894,11 @@ def create( ---------- store : StoreLike The array store that has already been initialized. - shape : ChunkCoords + shape : tuple[int, ...] The shape of the array. dtype : ZDTypeLike The data type of the array. - chunk_shape : ChunkCoords, optional + chunk_shape : tuple[int, ...], optional The shape of the Array's chunks. Zarr format 3 only. Zarr format 2 arrays should use `chunks` instead. If not specified, default are guessed based on the shape and dtype. @@ -1923,7 +1922,7 @@ def create( dimension_names : Iterable[str | None], optional The names of the dimensions (default is None). Zarr format 3 only. Zarr format 2 arrays should not use this parameter. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional The shape of the array's chunks. Zarr format 2 only. Zarr format 3 arrays should use ``chunk_shape`` instead. If not specified, default are guessed based on the shape and dtype. @@ -1990,13 +1989,13 @@ def _create( store: StoreLike, *, # v2 and v3 - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDTypeLike, zarr_format: ZarrFormat = 3, fill_value: Any | None = DEFAULT_FILL_VALUE, attributes: dict[str, JSON] | None = None, # v3 only - chunk_shape: ChunkCoords | None = None, + chunk_shape: tuple[int, ...] | None = None, chunk_key_encoding: ( ChunkKeyEncoding | tuple[Literal["default"], Literal[".", "/"]] @@ -2006,7 +2005,7 @@ def _create( codecs: Iterable[Codec | dict[str, JSON]] | None = None, dimension_names: DimensionNames = None, # v2 only - chunks: ChunkCoords | None = None, + chunks: tuple[int, ...] | None = None, dimension_separator: Literal[".", "/"] | None = None, order: MemoryOrder | None = None, filters: list[dict[str, JSON]] | None = None, @@ -2109,23 +2108,23 @@ def ndim(self) -> int: return self._async_array.ndim @property - def shape(self) -> ChunkCoords: + def shape(self) -> tuple[int, ...]: """Returns the shape of the array. Returns ------- - ChunkCoords + tuple[int, ...] The shape of the array. """ return self._async_array.shape @shape.setter - def shape(self, value: ChunkCoords) -> None: + def shape(self, value: tuple[int, ...]) -> None: """Sets the shape of the array by calling resize.""" self.resize(value) @property - def chunks(self) -> ChunkCoords: + def chunks(self) -> tuple[int, ...]: """Returns a tuple of integers describing the length of each dimension of a chunk of the array. If sharding is used the inner chunk shape is returned. @@ -2140,7 +2139,7 @@ def chunks(self) -> ChunkCoords: return self._async_array.chunks @property - def shards(self) -> ChunkCoords | None: + def shards(self) -> tuple[int, ...] | None: """Returns a tuple of integers describing the length of each dimension of a shard of the array. Returns None if sharding is not used. @@ -2262,7 +2261,7 @@ def compressors(self) -> tuple[numcodecs.abc.Codec, ...] | tuple[BytesBytesCodec return self._async_array.compressors @property - def cdata_shape(self) -> ChunkCoords: + def cdata_shape(self) -> tuple[int, ...]: """ The shape of the chunk grid for this array. """ @@ -2277,7 +2276,7 @@ def nchunks(self) -> int: def _iter_chunk_coords( self, origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None - ) -> Iterator[ChunkCoords]: + ) -> Iterator[tuple[int, ...]]: """ Create an iterator over the coordinates of chunks in chunk grid space. If the `origin` keyword is used, iteration will start at the chunk index specified by `origin`. @@ -2295,7 +2294,7 @@ def _iter_chunk_coords( Yields ------ - chunk_coords: ChunkCoords + chunk_coords: tuple[int, ...] The coordinates of each chunk in the selection. """ yield from self._async_array._iter_chunk_coords( @@ -3710,7 +3709,7 @@ def resize(self, new_shape: ShapeLike) -> None: """ sync(self._async_array.resize(new_shape)) - def append(self, data: npt.ArrayLike, axis: int = 0) -> ChunkCoords: + def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]: """Append `data` to `axis`. Parameters @@ -3922,11 +3921,11 @@ def _build_parents( class ShardsConfigParam(TypedDict): - shape: ChunkCoords + shape: tuple[int, ...] index_location: ShardingCodecIndexLocation | None -ShardsLike: TypeAlias = ChunkCoords | ShardsConfigParam | Literal["auto"] +ShardsLike: TypeAlias = tuple[int, ...] | ShardsConfigParam | Literal["auto"] async def from_array( @@ -3935,7 +3934,7 @@ async def from_array( data: Array | npt.ArrayLike, write_data: bool = True, name: str | None = None, - chunks: Literal["auto", "keep"] | ChunkCoords = "keep", + chunks: Literal["auto", "keep"] | tuple[int, ...] = "keep", shards: ShardsLike | None | Literal["keep"] = "keep", filters: FiltersLike | Literal["keep"] = "keep", compressors: CompressorsLike | Literal["keep"] = "keep", @@ -3965,22 +3964,22 @@ async def from_array( name : str or None, optional The name of the array within the store. If ``name`` is ``None``, the array will be located at the root of the store. - chunks : ChunkCoords or "auto" or "keep", optional + chunks : tuple[int, ...] or "auto" or "keep", optional Chunk shape of the array. Following values are supported: - "auto": Automatically determine the chunk shape based on the array's shape and dtype. - "keep": Retain the chunk shape of the data array if it is a zarr Array. - - ChunkCoords: A tuple of integers representing the chunk shape. + - tuple[int, ...]: A tuple of integers representing the chunk shape. If not specified, defaults to "keep" if data is a zarr Array, otherwise "auto". - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. Following values are supported: - "auto": Automatically determine the shard shape based on the array's shape and chunk shape. - "keep": Retain the shard shape of the data array if it is a zarr Array. - - ChunkCoords: A tuple of integers representing the shard shape. + - tuple[int, ...]: A tuple of integers representing the shard shape. - None: No sharding. If not specified, defaults to "keep" if data is a zarr Array, otherwise None. @@ -4168,7 +4167,9 @@ async def from_array( if write_data: if isinstance(data, Array): - async def _copy_array_region(chunk_coords: ChunkCoords | slice, _data: Array) -> None: + async def _copy_array_region( + chunk_coords: tuple[int, ...] | slice, _data: Array + ) -> None: arr = await _data._async_array.getitem(chunk_coords) await result.setitem(chunk_coords, arr) @@ -4197,7 +4198,7 @@ async def init_array( store_path: StorePath, shape: ShapeLike, dtype: ZDTypeLike, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -4217,14 +4218,14 @@ async def init_array( ---------- store_path : StorePath StorePath instance. The path attribute is the name of the array to initialize. - shape : ChunkCoords + shape : tuple[int, ...] Shape of the array. dtype : ZDTypeLike Data type of the array. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that @@ -4416,7 +4417,7 @@ async def create_array( shape: ShapeLike | None = None, dtype: ZDTypeLike | None = None, data: np.ndarray[Any, np.dtype[Any]] | None = None, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -4441,17 +4442,17 @@ async def create_array( name : str or None, optional The name of the array within the store. If ``name`` is ``None``, the array will be located at the root of the store. - shape : ChunkCoords, optional + shape : tuple[int, ...], optional Shape of the array. Can be ``None`` if ``data`` is provided. dtype : ZDTypeLike | None Data type of the array. Can be ``None`` if ``data`` is provided. data : Array-like data to use for initializing the array. If this parameter is provided, the ``shape`` and ``dtype`` parameters must be identical to ``data.shape`` and ``data.dtype``, or ``None``. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that @@ -4596,7 +4597,7 @@ async def create_array( def _parse_keep_array_attr( data: Array | npt.ArrayLike, - chunks: Literal["auto", "keep"] | ChunkCoords, + chunks: Literal["auto", "keep"] | tuple[int, ...], shards: ShardsLike | None | Literal["keep"], filters: FiltersLike | Literal["keep"], compressors: CompressorsLike | Literal["keep"], @@ -4607,7 +4608,7 @@ def _parse_keep_array_attr( chunk_key_encoding: ChunkKeyEncodingLike | None, dimension_names: DimensionNames, ) -> tuple[ - ChunkCoords | Literal["auto"], + tuple[int, ...] | Literal["auto"], ShardsLike | None, FiltersLike, CompressorsLike, diff --git a/src/zarr/core/array_spec.py b/src/zarr/core/array_spec.py index 279bf6edf0..c4dedaefea 100644 --- a/src/zarr/core/array_spec.py +++ b/src/zarr/core/array_spec.py @@ -16,7 +16,6 @@ from typing import NotRequired from zarr.core.buffer import BufferPrototype - from zarr.core.common import ChunkCoords from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar, ZDType @@ -88,7 +87,7 @@ def parse_array_config(data: ArrayConfigLike | None) -> ArrayConfig: @dataclass(frozen=True) class ArraySpec: - shape: ChunkCoords + shape: tuple[int, ...] dtype: ZDType[TBaseDType, TBaseScalar] fill_value: Any config: ArrayConfig @@ -96,7 +95,7 @@ class ArraySpec: def __init__( self, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TBaseDType, TBaseScalar], fill_value: Any, config: ArrayConfig, diff --git a/src/zarr/core/buffer/core.py b/src/zarr/core/buffer/core.py index 07bdb8c26e..d519b87d45 100644 --- a/src/zarr/core/buffer/core.py +++ b/src/zarr/core/buffer/core.py @@ -21,7 +21,7 @@ from typing import Self from zarr.codecs.bytes import Endian - from zarr.core.common import BytesLike, ChunkCoords + from zarr.core.common import BytesLike # Everything here is imported into ``zarr.core.buffer`` namespace. __all__: list[str] = [] @@ -59,7 +59,7 @@ def ndim(self) -> int: ... def size(self) -> int: ... @property - def shape(self) -> ChunkCoords: ... + def shape(self) -> tuple[int, ...]: ... def __len__(self) -> int: ... @@ -70,7 +70,7 @@ def __setitem__(self, key: slice, value: Any) -> None: ... def __array__(self) -> npt.NDArray[Any]: ... def reshape( - self, shape: ChunkCoords | Literal[-1], *, order: Literal["A", "C", "F"] = ... + self, shape: tuple[int, ...] | Literal[-1], *, order: Literal["A", "C", "F"] = ... ) -> Self: ... def view(self, dtype: npt.DTypeLike) -> Self: ... @@ -376,7 +376,7 @@ def create( @classmethod def empty( - cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" ) -> Self: """ Create an empty buffer with the given shape, dtype, and order. @@ -496,7 +496,7 @@ def byteorder(self) -> Endian: else: return Endian(sys.byteorder) - def reshape(self, newshape: ChunkCoords | Literal[-1]) -> Self: + def reshape(self, newshape: tuple[int, ...] | Literal[-1]) -> Self: return self.__class__(self._data.reshape(newshape)) def squeeze(self, axis: tuple[int, ...]) -> Self: diff --git a/src/zarr/core/buffer/cpu.py b/src/zarr/core/buffer/cpu.py index 9da0059d0b..34f92ece4a 100644 --- a/src/zarr/core/buffer/cpu.py +++ b/src/zarr/core/buffer/cpu.py @@ -20,7 +20,7 @@ from typing import Self from zarr.core.buffer.core import ArrayLike, NDArrayLike - from zarr.core.common import BytesLike, ChunkCoords + from zarr.core.common import BytesLike class Buffer(core.Buffer): @@ -162,7 +162,7 @@ def create( @classmethod def empty( - cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" ) -> Self: return cls(np.empty(shape=shape, dtype=dtype, order=order)) diff --git a/src/zarr/core/buffer/gpu.py b/src/zarr/core/buffer/gpu.py index 4eca197222..bfe977c50f 100644 --- a/src/zarr/core/buffer/gpu.py +++ b/src/zarr/core/buffer/gpu.py @@ -23,7 +23,7 @@ from collections.abc import Iterable from typing import Self - from zarr.core.common import BytesLike, ChunkCoords + from zarr.core.common import BytesLike try: import cupy as cp @@ -182,7 +182,7 @@ def create( @classmethod def empty( - cls, shape: ChunkCoords, dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" + cls, shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C" ) -> Self: return cls(cp.empty(shape=shape, dtype=dtype, order=order)) diff --git a/src/zarr/core/chunk_grids.py b/src/zarr/core/chunk_grids.py index 7fa1fc7e38..94c2e27674 100644 --- a/src/zarr/core/chunk_grids.py +++ b/src/zarr/core/chunk_grids.py @@ -15,8 +15,6 @@ from zarr.abc.metadata import Metadata from zarr.core.common import ( JSON, - ChunkCoords, - ChunkCoordsLike, ShapeLike, ceildiv, parse_named_configuration, @@ -32,13 +30,13 @@ def _guess_chunks( - shape: ShapeLike, + shape: tuple[int, ...] | int, typesize: int, *, increment_bytes: int = 256 * 1024, min_bytes: int = 128 * 1024, max_bytes: int = 64 * 1024 * 1024, -) -> ChunkCoords: +) -> tuple[int, ...]: """ Iteratively guess an appropriate chunk layout for an array, given its shape and the size of each element in bytes, and size constraints expressed in bytes. This logic is @@ -46,7 +44,7 @@ def _guess_chunks( Parameters ---------- - shape : ChunkCoords + shape : tuple[int, ...] The chunk shape. typesize : int The size, in bytes, of each element of the chunk. @@ -59,7 +57,7 @@ def _guess_chunks( Returns ------- - ChunkCoords + tuple[int, ...] """ if isinstance(shape, int): @@ -164,19 +162,19 @@ def from_dict(cls, data: dict[str, JSON] | ChunkGrid) -> ChunkGrid: raise ValueError(f"Unknown chunk grid. Got {name_parsed}.") @abstractmethod - def all_chunk_coords(self, array_shape: ChunkCoords) -> Iterator[ChunkCoords]: + def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: pass @abstractmethod - def get_nchunks(self, array_shape: ChunkCoords) -> int: + def get_nchunks(self, array_shape: tuple[int, ...]) -> int: pass @dataclass(frozen=True) class RegularChunkGrid(ChunkGrid): - chunk_shape: ChunkCoords + chunk_shape: tuple[int, ...] - def __init__(self, *, chunk_shape: ChunkCoordsLike) -> None: + def __init__(self, *, chunk_shape: ShapeLike) -> None: chunk_shape_parsed = parse_shapelike(chunk_shape) object.__setattr__(self, "chunk_shape", chunk_shape_parsed) @@ -190,12 +188,12 @@ def _from_dict(cls, data: dict[str, JSON]) -> Self: def to_dict(self) -> dict[str, JSON]: return {"name": "regular", "configuration": {"chunk_shape": tuple(self.chunk_shape)}} - def all_chunk_coords(self, array_shape: ChunkCoords) -> Iterator[ChunkCoords]: + def all_chunk_coords(self, array_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: return itertools.product( *(range(ceildiv(s, c)) for s, c in zip(array_shape, self.chunk_shape, strict=False)) ) - def get_nchunks(self, array_shape: ChunkCoords) -> int: + def get_nchunks(self, array_shape: tuple[int, ...]) -> int: return reduce( operator.mul, itertools.starmap(ceildiv, zip(array_shape, self.chunk_shape, strict=True)), diff --git a/src/zarr/core/chunk_key_encodings.py b/src/zarr/core/chunk_key_encodings.py index 91dfc90365..89a34e6052 100644 --- a/src/zarr/core/chunk_key_encodings.py +++ b/src/zarr/core/chunk_key_encodings.py @@ -10,7 +10,6 @@ from zarr.abc.metadata import Metadata from zarr.core.common import ( JSON, - ChunkCoords, parse_named_configuration, ) @@ -69,11 +68,11 @@ def to_dict(self) -> dict[str, JSON]: return {"name": self.name, "configuration": {"separator": self.separator}} @abstractmethod - def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: pass @abstractmethod - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: pass @@ -84,12 +83,12 @@ def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: class DefaultChunkKeyEncoding(ChunkKeyEncoding): name: Literal["default"] = "default" - def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: if chunk_key == "c": return () return tuple(map(int, chunk_key[1:].split(self.separator))) - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: return self.separator.join(map(str, ("c",) + chunk_coords)) @@ -97,9 +96,9 @@ def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: class V2ChunkKeyEncoding(ChunkKeyEncoding): name: Literal["v2"] = "v2" - def decode_chunk_key(self, chunk_key: str) -> ChunkCoords: + def decode_chunk_key(self, chunk_key: str) -> tuple[int, ...]: return tuple(map(int, chunk_key.split(self.separator))) - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: chunk_identifier = self.separator.join(map(str, chunk_coords)) return "0" if chunk_identifier == "" else chunk_identifier diff --git a/src/zarr/core/codec_pipeline.py b/src/zarr/core/codec_pipeline.py index 3bc3c1cfc7..63fcda7065 100644 --- a/src/zarr/core/codec_pipeline.py +++ b/src/zarr/core/codec_pipeline.py @@ -14,7 +14,7 @@ Codec, CodecPipeline, ) -from zarr.core.common import ChunkCoords, concurrent_map +from zarr.core.common import concurrent_map from zarr.core.config import config from zarr.core.indexing import SelectorTuple, is_scalar from zarr.errors import ZarrUserWarning @@ -134,7 +134,11 @@ def __iter__(self) -> Iterator[Codec]: yield from self.bytes_bytes_codecs def validate( - self, *, shape: ChunkCoords, dtype: ZDType[TBaseDType, TBaseScalar], chunk_grid: ChunkGrid + self, + *, + shape: tuple[int, ...], + dtype: ZDType[TBaseDType, TBaseScalar], + chunk_grid: ChunkGrid, ) -> None: for codec in self: codec.validate(shape=shape, dtype=dtype, chunk_grid=chunk_grid) diff --git a/src/zarr/core/common.py b/src/zarr/core/common.py index 4c0247426e..ed28fd2da4 100644 --- a/src/zarr/core/common.py +++ b/src/zarr/core/common.py @@ -36,9 +36,9 @@ ZMETADATA_V2_JSON = ".zmetadata" BytesLike = bytes | bytearray | memoryview -ShapeLike = tuple[int, ...] | int -ChunkCoords = tuple[int, ...] -ChunkCoordsLike = Iterable[int] +ShapeLike = Iterable[int] | int +# For backwards compatibility +ChunkCoords: tuple[int, ...] ZarrFormat = Literal[2, 3] NodeType = Literal["array", "group"] JSON = str | int | float | Mapping[str, "JSON"] | Sequence["JSON"] | None @@ -67,7 +67,7 @@ class NamedConfig(TypedDict, Generic[TName, TConfig]): """The configuration of the object.""" -def product(tup: ChunkCoords) -> int: +def product(tup: tuple[int, ...]) -> int: return functools.reduce(operator.mul, tup, 1) @@ -161,7 +161,7 @@ def parse_named_configuration( return name_parsed, configuration_parsed -def parse_shapelike(data: int | Iterable[int]) -> tuple[int, ...]: +def parse_shapelike(data: ShapeLike) -> tuple[int, ...]: if isinstance(data, int): if data < 0: raise ValueError(f"Expected a non-negative integer. Got {data} instead") diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 4bdc7b549f..0449e096f7 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -41,7 +41,6 @@ ZATTRS_JSON, ZGROUP_JSON, ZMETADATA_V2_JSON, - ChunkCoords, DimensionNames, NodeType, ShapeLike, @@ -1019,7 +1018,7 @@ async def create_array( shape: ShapeLike | None = None, dtype: ZDTypeLike | None = None, data: np.ndarray[Any, np.dtype[Any]] | None = None, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -1044,14 +1043,14 @@ async def create_array( name : str The name of the array relative to the group. If ``path`` is ``None``, the array will be located at the root of the store. - shape : ChunkCoords + shape : tuple[int, ...] Shape of the array. dtype : npt.DTypeLike Data type of the array. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that @@ -1197,7 +1196,7 @@ async def require_dataset( self, name: str, *, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: npt.DTypeLike = None, exact: bool = False, **kwargs: Any, @@ -1616,7 +1615,7 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An return await group_tree_async(self, max_depth=level) async def empty( - self, *, name: str, shape: ChunkCoords, **kwargs: Any + self, *, name: str, shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an empty array with the specified shape in this Group. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1639,7 +1638,7 @@ async def empty( return await async_api.empty(shape=shape, store=self.store_path, path=name, **kwargs) async def zeros( - self, *, name: str, shape: ChunkCoords, **kwargs: Any + self, *, name: str, shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with zero being used as the default value for uninitialized portions of the array. @@ -1660,7 +1659,7 @@ async def zeros( return await async_api.zeros(shape=shape, store=self.store_path, path=name, **kwargs) async def ones( - self, *, name: str, shape: ChunkCoords, **kwargs: Any + self, *, name: str, shape: tuple[int, ...], **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with one being used as the default value for uninitialized portions of the array. @@ -1681,7 +1680,7 @@ async def ones( return await async_api.ones(shape=shape, store=self.store_path, path=name, **kwargs) async def full( - self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any + self, *, name: str, shape: tuple[int, ...], fill_value: Any | None, **kwargs: Any ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. @@ -2450,7 +2449,7 @@ def create_array( shape: ShapeLike | None = None, dtype: ZDTypeLike | None = None, data: np.ndarray[Any, np.dtype[Any]] | None = None, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -2475,17 +2474,17 @@ def create_array( name : str The name of the array relative to the group. If ``path`` is ``None``, the array will be located at the root of the store. - shape : ChunkCoords, optional + shape : tuple[int, ...], optional Shape of the array. Can be ``None`` if ``data`` is provided. dtype : npt.DTypeLike | None Data type of the array. Can be ``None`` if ``data`` is provided. data : Array-like data to use for initializing the array. If this parameter is provided, the ``shape`` and ``dtype`` parameters must be identical to ``data.shape`` and ``data.dtype``, or ``None``. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that @@ -2660,7 +2659,7 @@ def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> Array: """ return Array(self._sync(self._async_group.require_array(name, shape=shape, **kwargs))) - def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: + def empty(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an empty array with the specified shape in this Group. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -2681,7 +2680,7 @@ def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: """ return Array(self._sync(self._async_group.empty(name=name, shape=shape, **kwargs))) - def zeros(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: + def zeros(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an array, with zero being used as the default value for uninitialized portions of the array. Parameters @@ -2700,7 +2699,7 @@ def zeros(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: """ return Array(self._sync(self._async_group.zeros(name=name, shape=shape, **kwargs))) - def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: + def ones(self, *, name: str, shape: tuple[int, ...], **kwargs: Any) -> Array: """Create an array, with one being used as the default value for uninitialized portions of the array. Parameters @@ -2720,7 +2719,7 @@ def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: return Array(self._sync(self._async_group.ones(name=name, shape=shape, **kwargs))) def full( - self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any + self, *, name: str, shape: tuple[int, ...], fill_value: Any | None, **kwargs: Any ) -> Array: """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. @@ -2846,8 +2845,8 @@ def array( *, shape: ShapeLike, dtype: npt.DTypeLike, - chunks: ChunkCoords | Literal["auto"] = "auto", - shards: ChunkCoords | Literal["auto"] | None = None, + chunks: tuple[int, ...] | Literal["auto"] = "auto", + shards: tuple[int, ...] | Literal["auto"] | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", compressor: CompressorLike = None, @@ -2874,14 +2873,14 @@ def array( name : str The name of the array relative to the group. If ``path`` is ``None``, the array will be located at the root of the store. - shape : ChunkCoords + shape : tuple[int, ...] Shape of the array. dtype : npt.DTypeLike Data type of the array. - chunks : ChunkCoords, optional + chunks : tuple[int, ...], optional Chunk shape of the array. If not specified, default are guessed based on the shape and dtype. - shards : ChunkCoords, optional + shards : tuple[int, ...], optional Shard shape of the array. The default value of ``None`` results in no sharding at all. filters : Iterable[Codec], optional Iterable of filters to apply to each chunk of the array, in order, before serializing that diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 15cf6f0f1a..16cb271601 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -33,7 +33,6 @@ from zarr.core.array import Array, AsyncArray from zarr.core.buffer import NDArrayLikeOrScalar from zarr.core.chunk_grids import ChunkGrid - from zarr.core.common import ChunkCoords IntSequence = list[int] | npt.NDArray[np.intp] @@ -75,7 +74,7 @@ class VindexInvalidSelectionError(IndexError): ) -def err_too_many_indices(selection: Any, shape: ChunkCoords) -> None: +def err_too_many_indices(selection: Any, shape: tuple[int, ...]) -> None: raise IndexError(f"too many indices for array; expected {len(shape)}, got {len(selection)}") @@ -90,8 +89,8 @@ def _zarr_array_to_int_or_bool_array(arr: Array) -> npt.NDArray[np.intp] | npt.N @runtime_checkable class Indexer(Protocol): - shape: ChunkCoords - drop_axes: ChunkCoords + shape: tuple[int, ...] + drop_axes: tuple[int, ...] def __iter__(self) -> Iterator[ChunkProjection]: ... @@ -105,7 +104,7 @@ def _iter_grid( origin: Sequence[int] | None = None, selection_shape: Sequence[int] | None = None, order: _ArrayIndexingOrder = "lexicographic", -) -> Iterator[ChunkCoords]: +) -> Iterator[tuple[int, ...]]: """ Iterate over the elements of grid of integers, with the option to restrict the domain of iteration to a contiguous subregion of that grid. @@ -283,7 +282,7 @@ def is_pure_orthogonal_indexing(selection: Selection, ndim: int) -> TypeGuard[Or ) -def get_chunk_shape(chunk_grid: ChunkGrid) -> ChunkCoords: +def get_chunk_shape(chunk_grid: ChunkGrid) -> tuple[int, ...]: from zarr.core.chunk_grids import RegularChunkGrid assert isinstance(chunk_grid, RegularChunkGrid), ( @@ -424,12 +423,12 @@ def __iter__(self) -> Iterator[ChunkDimProjection]: yield ChunkDimProjection(dim_chunk_ix, dim_chunk_sel, dim_out_sel, is_complete_chunk) -def check_selection_length(selection: SelectionNormalized, shape: ChunkCoords) -> None: +def check_selection_length(selection: SelectionNormalized, shape: tuple[int, ...]) -> None: if len(selection) > len(shape): err_too_many_indices(selection, shape) -def replace_ellipsis(selection: Any, shape: ChunkCoords) -> SelectionNormalized: +def replace_ellipsis(selection: Any, shape: tuple[int, ...]) -> SelectionNormalized: selection = ensure_tuple(selection) # count number of ellipsis present @@ -498,7 +497,7 @@ class ChunkProjection(NamedTuple): True if a complete chunk is indexed """ - chunk_coords: ChunkCoords + chunk_coords: tuple[int, ...] chunk_selection: tuple[Selector, ...] | npt.NDArray[np.intp] out_selection: tuple[Selector, ...] | npt.NDArray[np.intp] | slice is_complete_chunk: bool @@ -529,13 +528,13 @@ def is_basic_selection(selection: Any) -> TypeGuard[BasicSelection]: @dataclass(frozen=True) class BasicIndexer(Indexer): dim_indexers: list[IntDimIndexer | SliceDimIndexer] - shape: ChunkCoords - drop_axes: ChunkCoords + shape: tuple[int, ...] + drop_axes: tuple[int, ...] def __init__( self, selection: BasicSelection, - shape: ChunkCoords, + shape: tuple[int, ...], chunk_grid: ChunkGrid, ) -> None: chunk_shape = get_chunk_shape(chunk_grid) @@ -795,7 +794,7 @@ def slice_to_range(s: slice, length: int) -> range: return range(*s.indices(length)) -def ix_(selection: Any, shape: ChunkCoords) -> npt.NDArray[np.intp]: +def ix_(selection: Any, shape: tuple[int, ...]) -> npt.NDArray[np.intp]: """Convert an orthogonal selection to a numpy advanced (fancy) selection, like ``numpy.ix_`` but with support for slices and single ints.""" @@ -845,12 +844,12 @@ def oindex_set(a: npt.NDArray[Any], selection: Selection, value: Any) -> None: @dataclass(frozen=True) class OrthogonalIndexer(Indexer): dim_indexers: list[IntDimIndexer | SliceDimIndexer | IntArrayDimIndexer | BoolArrayDimIndexer] - shape: ChunkCoords - chunk_shape: ChunkCoords + shape: tuple[int, ...] + chunk_shape: tuple[int, ...] is_advanced: bool drop_axes: tuple[int, ...] - def __init__(self, selection: Selection, shape: ChunkCoords, chunk_grid: ChunkGrid) -> None: + def __init__(self, selection: Selection, shape: tuple[int, ...], chunk_grid: ChunkGrid) -> None: chunk_shape = get_chunk_shape(chunk_grid) # handle ellipsis @@ -979,11 +978,11 @@ async def getitem(self, selection: OrthogonalSelection | Array) -> NDArrayLikeOr @dataclass(frozen=True) class BlockIndexer(Indexer): dim_indexers: list[SliceDimIndexer] - shape: ChunkCoords - drop_axes: ChunkCoords + shape: tuple[int, ...] + drop_axes: tuple[int, ...] def __init__( - self, selection: BasicSelection, shape: ChunkCoords, chunk_grid: ChunkGrid + self, selection: BasicSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid ) -> None: chunk_shape = get_chunk_shape(chunk_grid) @@ -1078,7 +1077,7 @@ def __setitem__(self, selection: BasicSelection, value: npt.ArrayLike) -> None: def is_coordinate_selection( - selection: SelectionNormalized, shape: ChunkCoords + selection: SelectionNormalized, shape: tuple[int, ...] ) -> TypeGuard[CoordinateSelectionNormalized]: return ( isinstance(selection, tuple) @@ -1087,7 +1086,7 @@ def is_coordinate_selection( ) -def is_mask_selection(selection: Selection, shape: ChunkCoords) -> TypeGuard[MaskSelection]: +def is_mask_selection(selection: Selection, shape: tuple[int, ...]) -> TypeGuard[MaskSelection]: return ( isinstance(selection, tuple) and len(selection) == 1 @@ -1098,22 +1097,22 @@ def is_mask_selection(selection: Selection, shape: ChunkCoords) -> TypeGuard[Mas @dataclass(frozen=True) class CoordinateIndexer(Indexer): - sel_shape: ChunkCoords + sel_shape: tuple[int, ...] selection: CoordinateSelectionNormalized sel_sort: npt.NDArray[np.intp] | None chunk_nitems_cumsum: npt.NDArray[np.intp] chunk_rixs: npt.NDArray[np.intp] chunk_mixs: tuple[npt.NDArray[np.intp], ...] - shape: ChunkCoords - chunk_shape: ChunkCoords - drop_axes: ChunkCoords + shape: tuple[int, ...] + chunk_shape: tuple[int, ...] + drop_axes: tuple[int, ...] def __init__( - self, selection: CoordinateSelection, shape: ChunkCoords, chunk_grid: ChunkGrid + self, selection: CoordinateSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid ) -> None: chunk_shape = get_chunk_shape(chunk_grid) - cdata_shape: ChunkCoords + cdata_shape: tuple[int, ...] if shape == (): cdata_shape = (1,) else: @@ -1228,7 +1227,9 @@ def __iter__(self) -> Iterator[ChunkProjection]: @dataclass(frozen=True) class MaskIndexer(CoordinateIndexer): - def __init__(self, selection: MaskSelection, shape: ChunkCoords, chunk_grid: ChunkGrid) -> None: + def __init__( + self, selection: MaskSelection, shape: tuple[int, ...], chunk_grid: ChunkGrid + ) -> None: # some initial normalization selection_normalized = cast("tuple[MaskSelection]", ensure_tuple(selection)) selection_normalized = cast("tuple[MaskSelection]", replace_lists(selection_normalized)) @@ -1380,7 +1381,7 @@ def make_slice_selection(selection: Any) -> list[slice]: return ls -def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords: +def decode_morton(z: int, chunk_shape: tuple[int, ...]) -> tuple[int, ...]: # Inspired by compressed morton code as implemented in Neuroglancer # https://github.com/google/neuroglancer/blob/master/src/neuroglancer/datasource/precomputed/volume.md#compressed-morton-code bits = tuple(math.ceil(math.log2(c)) for c in chunk_shape) @@ -1398,9 +1399,9 @@ def decode_morton(z: int, chunk_shape: ChunkCoords) -> ChunkCoords: return tuple(out) -def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]: +def morton_order_iter(chunk_shape: tuple[int, ...]) -> Iterator[tuple[int, ...]]: i = 0 - order: list[ChunkCoords] = [] + order: list[tuple[int, ...]] = [] while len(order) < product(chunk_shape): m = decode_morton(i, chunk_shape) if m not in order and all(x < y for x, y in zip(m, chunk_shape, strict=False)): @@ -1410,12 +1411,12 @@ def morton_order_iter(chunk_shape: ChunkCoords) -> Iterator[ChunkCoords]: yield order[j] -def c_order_iter(chunks_per_shard: ChunkCoords) -> Iterator[ChunkCoords]: +def c_order_iter(chunks_per_shard: tuple[int, ...]) -> Iterator[tuple[int, ...]]: return itertools.product(*(range(x) for x in chunks_per_shard)) def get_indexer( - selection: SelectionWithFields, shape: ChunkCoords, chunk_grid: ChunkGrid + selection: SelectionWithFields, shape: tuple[int, ...], chunk_grid: ChunkGrid ) -> Indexer: _, pure_selection = pop_fields(selection) if is_pure_fancy_indexing(pure_selection, len(shape)): diff --git a/src/zarr/core/metadata/v2.py b/src/zarr/core/metadata/v2.py index 9ad6b3bc42..56c7b77df7 100644 --- a/src/zarr/core/metadata/v2.py +++ b/src/zarr/core/metadata/v2.py @@ -19,7 +19,6 @@ import numpy.typing as npt from zarr.core.buffer import Buffer, BufferPrototype - from zarr.core.common import ChunkCoords from zarr.core.dtype.wrapper import ( TBaseDType, TBaseScalar, @@ -62,8 +61,8 @@ class ArrayV2MetadataDict(TypedDict): @dataclass(frozen=True, kw_only=True) class ArrayV2Metadata(Metadata): - shape: ChunkCoords - chunks: ChunkCoords + shape: tuple[int, ...] + chunks: tuple[int, ...] dtype: ZDType[TBaseDType, TBaseScalar] fill_value: int | float | str | bytes | None = None order: MemoryOrder = "C" @@ -76,9 +75,9 @@ class ArrayV2Metadata(Metadata): def __init__( self, *, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: ZDType[TDType_co, TScalar_co], - chunks: ChunkCoords, + chunks: tuple[int, ...], fill_value: Any, order: MemoryOrder, dimension_separator: Literal[".", "/"] = ".", @@ -124,7 +123,7 @@ def chunk_grid(self) -> RegularChunkGrid: return RegularChunkGrid(chunk_shape=self.chunks) @property - def shards(self) -> ChunkCoords | None: + def shards(self) -> tuple[int, ...] | None: return None def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: @@ -230,7 +229,7 @@ def to_dict(self) -> dict[str, JSON]: return zarray_dict def get_chunk_spec( - self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype + self, _chunk_coords: tuple[int, ...], array_config: ArrayConfig, prototype: BufferPrototype ) -> ArraySpec: return ArraySpec( shape=self.chunks, @@ -240,11 +239,11 @@ def get_chunk_spec( prototype=prototype, ) - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: chunk_identifier = self.dimension_separator.join(map(str, chunk_coords)) return "0" if chunk_identifier == "" else chunk_identifier - def update_shape(self, shape: ChunkCoords) -> Self: + def update_shape(self, shape: tuple[int, ...]) -> Self: return replace(self, shape=shape) def update_attributes(self, attributes: dict[str, JSON]) -> Self: diff --git a/src/zarr/core/metadata/v3.py b/src/zarr/core/metadata/v3.py index 6f79fb4b09..e17edb999c 100644 --- a/src/zarr/core/metadata/v3.py +++ b/src/zarr/core/metadata/v3.py @@ -12,7 +12,7 @@ from zarr.core.buffer import Buffer, BufferPrototype from zarr.core.chunk_grids import ChunkGrid - from zarr.core.common import JSON, ChunkCoords + from zarr.core.common import JSON from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar @@ -28,7 +28,6 @@ from zarr.core.common import ( JSON, ZARR_JSON, - ChunkCoords, DimensionNames, parse_named_configuration, parse_shapelike, @@ -138,7 +137,7 @@ class ArrayV3MetadataDict(TypedDict): @dataclass(frozen=True, kw_only=True) class ArrayV3Metadata(Metadata): - shape: ChunkCoords + shape: tuple[int, ...] data_type: ZDType[TBaseDType, TBaseScalar] chunk_grid: ChunkGrid chunk_key_encoding: ChunkKeyEncoding @@ -224,7 +223,7 @@ def dtype(self) -> ZDType[TBaseDType, TBaseScalar]: return self.data_type @property - def chunks(self) -> ChunkCoords: + def chunks(self) -> tuple[int, ...]: if isinstance(self.chunk_grid, RegularChunkGrid): from zarr.codecs.sharding import ShardingCodec @@ -242,7 +241,7 @@ def chunks(self) -> ChunkCoords: raise NotImplementedError(msg) @property - def shards(self) -> ChunkCoords | None: + def shards(self) -> tuple[int, ...] | None: if isinstance(self.chunk_grid, RegularChunkGrid): from zarr.codecs.sharding import ShardingCodec @@ -267,7 +266,7 @@ def inner_codecs(self) -> tuple[Codec, ...]: return self.codecs def get_chunk_spec( - self, _chunk_coords: ChunkCoords, array_config: ArrayConfig, prototype: BufferPrototype + self, _chunk_coords: tuple[int, ...], array_config: ArrayConfig, prototype: BufferPrototype ) -> ArraySpec: assert isinstance(self.chunk_grid, RegularChunkGrid), ( "Currently, only regular chunk grid is supported" @@ -280,7 +279,7 @@ def get_chunk_spec( prototype=prototype, ) - def encode_chunk_key(self, chunk_coords: ChunkCoords) -> str: + def encode_chunk_key(self, chunk_coords: tuple[int, ...]) -> str: return self.chunk_key_encoding.encode_chunk_key(chunk_coords) def to_buffer_dict(self, prototype: BufferPrototype) -> dict[str, Buffer]: @@ -345,7 +344,7 @@ def to_dict(self) -> dict[str, JSON]: return out_dict - def update_shape(self, shape: ChunkCoords) -> Self: + def update_shape(self, shape: tuple[int, ...]) -> Self: return replace(self, shape=shape) def update_attributes(self, attributes: dict[str, JSON]) -> Self: diff --git a/src/zarr/testing/buffer.py b/src/zarr/testing/buffer.py index 8cbfb2414a..1e167b2156 100644 --- a/src/zarr/testing/buffer.py +++ b/src/zarr/testing/buffer.py @@ -13,8 +13,6 @@ from collections.abc import Iterable from typing import Self - from zarr.core.common import ChunkCoords - __all__ = [ "NDBufferUsingTestNDArrayLike", @@ -56,7 +54,7 @@ def create( @classmethod def empty( cls, - shape: ChunkCoords, + shape: tuple[int, ...], dtype: npt.DTypeLike, order: Literal["C", "F"] = "C", ) -> Self: diff --git a/tests/conftest.py b/tests/conftest.py index a1bf423c06..839be34e01 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,7 +21,14 @@ _parse_chunk_key_encoding, ) from zarr.core.chunk_grids import RegularChunkGrid, _auto_partition -from zarr.core.common import JSON, DimensionNames, parse_shapelike +from zarr.core.common import ( + JSON, + DimensionNames, + MemoryOrder, + ShapeLike, + ZarrFormat, + parse_shapelike, +) from zarr.core.config import config as zarr_config from zarr.core.dtype import ( get_data_type_from_native_dtype, @@ -40,8 +47,10 @@ from zarr.abc.codec import Codec 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.chunk_key_encodings import ( + ChunkKeyEncoding, + ChunkKeyEncodingLike, + ) from zarr.core.dtype.wrapper import ZDType @@ -152,7 +161,7 @@ def reset_config() -> Generator[None, None, None]: @dataclass class ArrayRequest: - shape: ChunkCoords + shape: tuple[int, ...] dtype: str order: MemoryOrder @@ -229,7 +238,7 @@ def create_array_metadata( *, shape: ShapeLike, dtype: npt.DTypeLike, - chunks: ChunkCoords | Literal["auto"], + chunks: tuple[int, ...] | Literal["auto"], shards: None, filters: FiltersLike, compressors: CompressorsLike, @@ -248,7 +257,7 @@ def create_array_metadata( *, shape: ShapeLike, dtype: npt.DTypeLike, - chunks: ChunkCoords | Literal["auto"], + chunks: tuple[int, ...] | Literal["auto"], shards: ShardsLike | None, filters: FiltersLike, compressors: CompressorsLike, @@ -267,7 +276,7 @@ def create_array_metadata( *, shape: ShapeLike, dtype: npt.DTypeLike, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", @@ -369,7 +378,7 @@ def create_array_metadata( @overload def meta_from_array( array: np.ndarray[Any, Any], - chunks: ChunkCoords | Literal["auto"], + chunks: tuple[int, ...] | Literal["auto"], shards: None, filters: FiltersLike, compressors: CompressorsLike, @@ -386,7 +395,7 @@ def meta_from_array( @overload def meta_from_array( array: np.ndarray[Any, Any], - chunks: ChunkCoords | Literal["auto"], + chunks: tuple[int, ...] | Literal["auto"], shards: ShardsLike | None, filters: FiltersLike, compressors: CompressorsLike, @@ -405,7 +414,7 @@ def meta_from_array( def meta_from_array( array: np.ndarray[Any, Any], *, - chunks: ChunkCoords | Literal["auto"] = "auto", + chunks: tuple[int, ...] | Literal["auto"] = "auto", shards: ShardsLike | None = None, filters: FiltersLike = "auto", compressors: CompressorsLike = "auto", diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py index a2dad41a1b..dfedbb83de 100644 --- a/tests/test_codecs/test_codecs.py +++ b/tests/test_codecs/test_codecs.py @@ -28,7 +28,7 @@ from zarr.abc.codec import Codec from zarr.abc.store import Store from zarr.core.buffer.core import NDArrayLikeOrScalar - from zarr.core.common import ChunkCoords, MemoryOrder + from zarr.core.common import MemoryOrder @dataclass(frozen=True) @@ -215,7 +215,7 @@ def test_morton() -> None: [3, 2, 1, 6, 4, 5, 2], ], ) -def test_morton2(shape: ChunkCoords) -> None: +def test_morton2(shape: tuple[int, ...]) -> None: order = list(morton_order_iter(shape)) for i, x in enumerate(order): assert x not in order[:i] # no duplicates diff --git a/tests/test_indexing.py b/tests/test_indexing.py index 24b4b65505..a78fe3c706 100644 --- a/tests/test_indexing.py +++ b/tests/test_indexing.py @@ -33,7 +33,6 @@ from zarr.core.buffer import BufferPrototype from zarr.core.buffer.core import Buffer - from zarr.core.common import ChunkCoords @pytest.fixture @@ -44,7 +43,7 @@ async def store() -> AsyncGenerator[StorePath]: def zarr_array_from_numpy_array( store: StorePath, a: npt.NDArray[Any], - chunk_shape: ChunkCoords | None = None, + chunk_shape: tuple[int, ...] | None = None, ) -> zarr.Array: z = zarr.create_array( store=store / str(uuid4()),