diff --git a/changes/3304.feature.rst b/changes/3304.feature.rst new file mode 100644 index 0000000000..45dbd85731 --- /dev/null +++ b/changes/3304.feature.rst @@ -0,0 +1,2 @@ +The `Array` class can now also be parametrized in the same manner as the `AsyncArray` class, allowing Zarr format v2 and v3 `Array`s to be distinguished. +New types have been added to `zarr.types` to help with this. diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 78b68caf73..e909f2e40d 100644 --- a/src/zarr/api/asynchronous.py +++ b/src/zarr/api/asynchronous.py @@ -3,7 +3,7 @@ import asyncio import dataclasses import warnings -from typing import TYPE_CHECKING, Any, Literal, cast +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast import numpy as np import numpy.typing as npt @@ -38,7 +38,7 @@ GroupMetadata, create_hierarchy, ) -from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata from zarr.errors import ( GroupNotFoundError, NodeTypeValidationError, @@ -58,9 +58,10 @@ from zarr.core.buffer import NDArrayLikeOrScalar from zarr.core.chunk_key_encodings import ChunkKeyEncoding from zarr.storage import StoreLike + from zarr.types import AnyArray, AnyAsyncArray # TODO: this type could use some more thought - ArrayLike = AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array | npt.NDArray[Any] + ArrayLike: TypeAlias = AnyAsyncArray | AnyArray | npt.NDArray[Any] PathLike = str __all__ = [ @@ -312,7 +313,7 @@ async def open( path: str | None = None, storage_options: dict[str, Any] | None = None, **kwargs: Any, # TODO: type kwargs as valid args to open_array -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup: +) -> AnyAsyncArray | AsyncGroup: """Convenience function to open a group or array using file-mode-like semantics. Parameters @@ -567,9 +568,7 @@ async def tree(grp: AsyncGroup, expand: bool | None = None, level: int | None = return await grp.tree(expand=expand, level=level) -async def array( - data: npt.ArrayLike | Array, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def array(data: npt.ArrayLike | AnyArray, **kwargs: Any) -> AnyAsyncArray: """Create an array filled with `data`. Parameters @@ -901,7 +900,7 @@ async def create( storage_options: dict[str, Any] | None = None, config: ArrayConfigLike | None = None, **kwargs: Any, -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> AnyAsyncArray: """Create an array. Parameters @@ -1073,9 +1072,7 @@ async def create( ) -async def empty( - shape: ChunkCoords, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def empty(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """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. @@ -1096,9 +1093,7 @@ async def empty( return await create(shape=shape, fill_value=None, **kwargs) -async def empty_like( - a: ArrayLike, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def empty_like(a: ArrayLike, **kwargs: Any) -> AnyAsyncArray: """Create an empty array like `a`. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1125,9 +1120,7 @@ async def empty_like( # TODO: add type annotations for fill_value and kwargs -async def full( - shape: ChunkCoords, fill_value: Any, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AnyAsyncArray: """Create an array, with `fill_value` being used as the default value for uninitialized portions of the array. @@ -1149,9 +1142,7 @@ async def full( # TODO: add type annotations for kwargs -async def full_like( - a: ArrayLike, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def full_like(a: ArrayLike, **kwargs: Any) -> AnyAsyncArray: """Create a filled array like `a`. Parameters @@ -1172,9 +1163,7 @@ async def full_like( return await full(**like_kwargs) -async def ones( - shape: ChunkCoords, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def ones(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """Create an array, with one being used as the default value for uninitialized portions of the array. @@ -1193,9 +1182,7 @@ async def ones( return await create(shape=shape, fill_value=1, **kwargs) -async def ones_like( - a: ArrayLike, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def ones_like(a: ArrayLike, **kwargs: Any) -> AnyAsyncArray: """Create an array of ones like `a`. Parameters @@ -1222,7 +1209,7 @@ async def open_array( path: PathLike = "", storage_options: dict[str, Any] | None = None, **kwargs: Any, # TODO: type kwargs as valid args to save -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> AnyAsyncArray: """Open an array using file-mode-like semantics. Parameters @@ -1270,9 +1257,7 @@ async def open_array( raise -async def open_like( - a: ArrayLike, path: str, **kwargs: Any -) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: +async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AnyAsyncArray: """Open a persistent array like `a`. Parameters @@ -1295,9 +1280,7 @@ async def open_like( return await open_array(path=path, **like_kwargs) -async def zeros( - shape: ChunkCoords, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def zeros(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """Create an array, with zero being used as the default value for uninitialized portions of the array. @@ -1316,9 +1299,7 @@ async def zeros( return await create(shape=shape, fill_value=0, **kwargs) -async def zeros_like( - a: ArrayLike, **kwargs: Any -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +async def zeros_like(a: ArrayLike, **kwargs: Any) -> AnyAsyncArray: """Create an array of zeros like `a`. Parameters diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index ed1ae2cf2a..1b71d3475b 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -41,6 +41,7 @@ ) from zarr.core.dtype import ZDTypeLike from zarr.storage import StoreLike + from zarr.types import AnyArray __all__ = [ "array", @@ -169,7 +170,7 @@ def open( path: str | None = None, storage_options: dict[str, Any] | None = None, **kwargs: Any, # TODO: type kwargs as valid args to async_api.open -) -> Array | Group: +) -> AnyArray | Group: """Open a group or array using file-mode-like semantics. Parameters @@ -366,7 +367,7 @@ def tree(grp: Group, expand: bool | None = None, level: int | None = None) -> An # TODO: add type annotations for kwargs -def array(data: npt.ArrayLike | Array, **kwargs: Any) -> Array: +def array(data: npt.ArrayLike | AnyArray, **kwargs: Any) -> AnyArray: """Create an array filled with `data`. Parameters @@ -634,7 +635,7 @@ def create( storage_options: dict[str, Any] | None = None, config: ArrayConfigLike | None = None, **kwargs: Any, -) -> Array: +) -> AnyArray: """Create an array. Parameters @@ -770,7 +771,7 @@ def create_array( overwrite: bool = False, config: ArrayConfigLike | None = None, write_data: bool = True, -) -> Array: +) -> AnyArray: """Create an array. This function wraps :func:`zarr.core.array.create_array`. @@ -918,7 +919,7 @@ def create_array( def from_array( store: str | StoreLike, *, - data: Array | npt.ArrayLike, + data: AnyArray | npt.ArrayLike, write_data: bool = True, name: str | None = None, chunks: Literal["auto", "keep"] | ChunkCoords = "keep", @@ -935,7 +936,7 @@ def from_array( storage_options: dict[str, Any] | None = None, overwrite: bool = False, config: ArrayConfigLike | None = None, -) -> Array: +) -> AnyArray: """Create an array from an existing array or array-like. Parameters @@ -1129,7 +1130,7 @@ def from_array( # TODO: add type annotations for kwargs -def empty(shape: ChunkCoords, **kwargs: Any) -> Array: +def empty(shape: ChunkCoords, **kwargs: Any) -> AnyArray: """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. @@ -1156,7 +1157,7 @@ def empty(shape: ChunkCoords, **kwargs: Any) -> Array: # TODO: move ArrayLike to common module # TODO: add type annotations for kwargs -def empty_like(a: ArrayLike, **kwargs: Any) -> Array: +def empty_like(a: ArrayLike, **kwargs: Any) -> AnyArray: """Create an empty array like another array. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1182,7 +1183,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: ChunkCoords, fill_value: Any, **kwargs: Any) -> AnyArray: """Create an array with a default fill value. Parameters @@ -1204,7 +1205,7 @@ def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> Array: # TODO: move ArrayLike to common module # TODO: add type annotations for kwargs -def full_like(a: ArrayLike, **kwargs: Any) -> Array: +def full_like(a: ArrayLike, **kwargs: Any) -> AnyArray: """Create a filled array like another array. Parameters @@ -1223,7 +1224,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: ChunkCoords, **kwargs: Any) -> AnyArray: """Create an array with a fill value of one. Parameters @@ -1242,7 +1243,7 @@ def ones(shape: ChunkCoords, **kwargs: Any) -> Array: # TODO: add type annotations for kwargs -def ones_like(a: ArrayLike, **kwargs: Any) -> Array: +def ones_like(a: ArrayLike, **kwargs: Any) -> AnyArray: """Create an array of ones like another array. Parameters @@ -1268,7 +1269,7 @@ def open_array( path: PathLike = "", storage_options: dict[str, Any] | None = None, **kwargs: Any, -) -> Array: +) -> AnyArray: """Open an array using file-mode-like semantics. Parameters @@ -1304,7 +1305,7 @@ def open_array( # TODO: add type annotations for kwargs -def open_like(a: ArrayLike, path: str, **kwargs: Any) -> Array: +def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AnyArray: """Open a persistent array like another array. Parameters @@ -1325,7 +1326,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: ChunkCoords, **kwargs: Any) -> AnyArray: """Create an array with a fill value of zero. Parameters @@ -1344,7 +1345,7 @@ def zeros(shape: ChunkCoords, **kwargs: Any) -> Array: # TODO: add type annotations for kwargs -def zeros_like(a: ArrayLike, **kwargs: Any) -> Array: +def zeros_like(a: ArrayLike, **kwargs: Any) -> AnyArray: """Create an array of zeros like another array. Parameters diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 311a0eb986..1655243fed 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -140,6 +140,7 @@ from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.core.group import AsyncGroup from zarr.storage import StoreLike + from zarr.types import AnyArray, AnyAsyncArray, AsyncArrayV2, AsyncArrayV3 # Array and AsyncArray are defined in the base ``zarr`` namespace @@ -293,7 +294,7 @@ class AsyncArray(Generic[T_ArrayMetadata]): @overload def __init__( - self: AsyncArray[ArrayV2Metadata], + self: AsyncArrayV2, metadata: ArrayV2Metadata | ArrayV2MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, @@ -301,7 +302,7 @@ def __init__( @overload def __init__( - self: AsyncArray[ArrayV3Metadata], + self: AsyncArrayV3, metadata: ArrayV3Metadata | ArrayV3MetadataDict, store_path: StorePath, config: ArrayConfigLike | None = None, @@ -347,7 +348,7 @@ async def create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV2Metadata]: ... + ) -> AsyncArrayV2: ... # this overload defines the function signature when zarr_format is 3 @overload @@ -376,7 +377,7 @@ async def create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV3Metadata]: ... + ) -> AsyncArrayV3: ... @overload @classmethod @@ -404,7 +405,7 @@ async def create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV3Metadata]: ... + ) -> AsyncArrayV3: ... @overload @classmethod @@ -438,7 +439,7 @@ async def create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: ... + ) -> AnyAsyncArray: ... @classmethod @deprecated("Use zarr.api.asynchronous.create_array instead.", category=ZarrDeprecationWarning) @@ -472,7 +473,7 @@ async def create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Method to create a new asynchronous array instance. .. deprecated:: 3.0.0 @@ -613,7 +614,7 @@ async def _create( overwrite: bool = False, data: npt.ArrayLike | None = None, config: ArrayConfigLike | None = None, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Method to create a new asynchronous array instance. See :func:`AsyncArray.create` for more details. Deprecated in favor of :func:`zarr.api.asynchronous.create_array`. @@ -635,7 +636,7 @@ async def _create( _chunks = normalize_chunks(chunk_shape, shape, item_size) config_parsed = parse_array_config(config) - result: AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] + result: AnyAsyncArray if zarr_format == 3: if dimension_separator is not None: raise ValueError( @@ -779,7 +780,7 @@ async def _create_v3( dimension_names: DimensionNames = None, attributes: dict[str, JSON] | None = None, overwrite: bool = False, - ) -> AsyncArray[ArrayV3Metadata]: + ) -> AsyncArrayV3: if overwrite: if store_path.store.supports_deletes: await store_path.delete_dir() @@ -860,7 +861,7 @@ async def _create_v2( compressor: CompressorLike = "auto", attributes: dict[str, JSON] | None = None, overwrite: bool = False, - ) -> AsyncArray[ArrayV2Metadata]: + ) -> AsyncArrayV2: if overwrite: if store_path.store.supports_deletes: await store_path.delete_dir() @@ -904,7 +905,7 @@ def from_dict( cls, store_path: StorePath, data: dict[str, JSON], - ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: + ) -> AnyAsyncArray: """ Create a Zarr array from a dictionary, with support for both Zarr format 2 and 3 metadata. @@ -920,7 +921,7 @@ def from_dict( Returns ------- - AsyncArray[ArrayV3Metadata] or AsyncArray[ArrayV2Metadata] + AsyncArrayV3 or AsyncArrayV2 The created Zarr array, either using Zarr format 2 or 3 metadata based on the provided data. Raises @@ -936,7 +937,7 @@ async def open( cls, store: StoreLike, zarr_format: ZarrFormat | None = 3, - ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: + ) -> AnyAsyncArray: """ Async method to open an existing Zarr array from a given store. @@ -1847,12 +1848,12 @@ def _info( # TODO: Array can be a frozen data class again once property setters (e.g. shape) are removed @dataclass(frozen=False) -class Array: +class Array(Generic[T_ArrayMetadata]): """ A Zarr array. """ - _async_array: AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] + _async_array: AsyncArray[T_ArrayMetadata] @classmethod @deprecated("Use zarr.create_array instead.", category=ZarrDeprecationWarning) @@ -1885,7 +1886,7 @@ def create( # runtime overwrite: bool = False, config: ArrayConfigLike | None = None, - ) -> Array: + ) -> AnyArray: """Creates a new Array instance from an initialized store. .. deprecated:: 3.0.0 @@ -2014,7 +2015,7 @@ def _create( # runtime overwrite: bool = False, config: ArrayConfigLike | None = None, - ) -> Array: + ) -> AnyArray: """Creates a new Array instance from an initialized store. See :func:`Array.create` for more details. Deprecated in favor of :func:`zarr.create_array`. @@ -2040,14 +2041,14 @@ def _create( config=config, ), ) - return cls(async_array) + return Array(async_array) @classmethod def from_dict( cls, store_path: StorePath, data: dict[str, JSON], - ) -> Array: + ) -> AnyArray: """ Create a Zarr array from a dictionary. @@ -2071,13 +2072,13 @@ def from_dict( If the dictionary data is invalid or missing required fields for array creation. """ async_array = AsyncArray.from_dict(store_path=store_path, data=data) - return cls(async_array) + return Array(async_array) @classmethod def open( cls, store: StoreLike, - ) -> Array: + ) -> AnyArray: """Opens an existing Array from a store. Parameters @@ -2091,7 +2092,7 @@ def open( Array opened from the store. """ async_array = sync(AsyncArray.open(store)) - return cls(async_array) + return Array(async_array) @property def store(self) -> Store: @@ -3746,7 +3747,7 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> ChunkCoords: """ return sync(self._async_array.append(data, axis=axis)) - def update_attributes(self, new_attributes: dict[str, JSON]) -> Array: + def update_attributes(self, new_attributes: dict[str, JSON]) -> Self: """ Update the array's attributes. @@ -3771,11 +3772,8 @@ def update_attributes(self, new_attributes: dict[str, JSON]) -> Array: - The updated attributes will be merged with existing attributes, and any conflicts will be overwritten by the new values. """ - # TODO: remove this cast when type inference improves new_array = sync(self._async_array.update_attributes(new_attributes)) - # TODO: remove this cast when type inference improves - _new_array = cast("AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]", new_array) - return type(self)(_new_array) + return type(self)(new_array) def __repr__(self) -> str: return f"" @@ -3835,7 +3833,7 @@ def info_complete(self) -> Any: async def chunks_initialized( - array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], + array: AnyAsyncArray, ) -> tuple[str, ...]: """ Return the keys of the chunks that have been persisted to the storage backend. @@ -3867,7 +3865,7 @@ async def chunks_initialized( def _build_parents( - node: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup, + node: AnyAsyncArray | AsyncGroup, ) -> list[AsyncGroup]: from zarr.core.group import AsyncGroup, GroupMetadata @@ -3932,7 +3930,7 @@ class ShardsConfigParam(TypedDict): async def from_array( store: str | StoreLike, *, - data: Array | npt.ArrayLike, + data: AnyArray | npt.ArrayLike, write_data: bool = True, name: str | None = None, chunks: Literal["auto", "keep"] | ChunkCoords = "keep", @@ -3949,7 +3947,7 @@ async def from_array( storage_options: dict[str, Any] | None = None, overwrite: bool = False, config: ArrayConfig | ArrayConfigLike | None = None, -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> AnyAsyncArray: """Create an array from an existing array or array-like. Parameters @@ -4168,7 +4166,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: ChunkCoords | slice, _data: AnyArray + ) -> None: arr = await _data._async_array.getitem(chunk_coords) await result.setitem(chunk_coords, arr) @@ -4210,7 +4210,7 @@ async def init_array( dimension_names: DimensionNames = None, overwrite: bool = False, config: ArrayConfigLike | None, -) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]: +) -> AnyAsyncArray: """Create and persist an array metadata document. Parameters @@ -4431,7 +4431,7 @@ async def create_array( overwrite: bool = False, config: ArrayConfigLike | None = None, write_data: bool = True, -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> AnyAsyncArray: """Create an array. Parameters @@ -4595,7 +4595,7 @@ async def create_array( def _parse_keep_array_attr( - data: Array | npt.ArrayLike, + data: AnyArray | npt.ArrayLike, chunks: Literal["auto", "keep"] | ChunkCoords, shards: ShardsLike | None | Literal["keep"], filters: FiltersLike | Literal["keep"], diff --git a/src/zarr/core/attributes.py b/src/zarr/core/attributes.py index e699c4f66d..a4fe80020f 100644 --- a/src/zarr/core/attributes.py +++ b/src/zarr/core/attributes.py @@ -8,12 +8,12 @@ if TYPE_CHECKING: from collections.abc import Iterator - from zarr.core.array import Array from zarr.core.group import Group + from zarr.types import AnyArray class Attributes(MutableMapping[str, JSON]): - def __init__(self, obj: Array | Group) -> None: + def __init__(self, obj: AnyArray | Group) -> None: # key=".zattrs", read_only=False, cache=True, synchronizer=None self._obj = obj diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index 4bdc7b549f..7d0dc9ccf5 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -79,6 +79,7 @@ from zarr.core.chunk_key_encodings import ChunkKeyEncodingLike from zarr.core.common import MemoryOrder from zarr.core.dtype import ZDTypeLike + from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 logger = logging.getLogger("zarr.group") @@ -112,7 +113,11 @@ def parse_attributes(data: Any) -> dict[str, Any]: @overload -def _parse_async_node(node: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]) -> Array: ... +def _parse_async_node(node: AsyncArrayV3) -> ArrayV3: ... + + +@overload +def _parse_async_node(node: AsyncArrayV2) -> ArrayV2: ... @overload @@ -120,8 +125,8 @@ def _parse_async_node(node: AsyncGroup) -> Group: ... def _parse_async_node( - node: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup, -) -> Array | Group: + node: AnyAsyncArray | AsyncGroup, +) -> AnyArray | Group: """Wrap an AsyncArray in an Array, or an AsyncGroup in a Group.""" if isinstance(node, AsyncArray): return Array(node) @@ -697,7 +702,7 @@ async def setitem(self, key: str, value: Any) -> None: async def getitem( self, key: str, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup: + ) -> AnyAsyncArray | AsyncGroup: """ Get a subarray or subgroup from the group. @@ -725,7 +730,7 @@ async def getitem( def _getitem_consolidated( self, store_path: StorePath, key: str, prefix: str - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup: + ) -> AnyAsyncArray | AsyncGroup: # getitem, in the special case where we have consolidated metadata. # Note that this is a regular def (non async) function. # This shouldn't do any additional I/O. @@ -788,7 +793,7 @@ async def delitem(self, key: str) -> None: async def get( self, key: str, default: DefaultT | None = None - ) -> AsyncArray[Any] | AsyncGroup | DefaultT | None: + ) -> AnyAsyncArray | AsyncGroup | DefaultT | None: """Obtain a group member, returning default if not found. Parameters @@ -983,9 +988,7 @@ async def require_group(self, name: str, overwrite: bool = False) -> AsyncGroup: grp = await self.create_group(name, overwrite=True) else: try: - item: ( - AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] - ) = await self.getitem(name) + item: AsyncGroup | AnyAsyncArray = await self.getitem(name) if not isinstance(item, AsyncGroup): raise TypeError( f"Incompatible object ({item.__class__.__name__}) already exists" @@ -1034,7 +1037,7 @@ async def create_array( overwrite: bool = False, config: ArrayConfigLike | None = None, write_data: bool = True, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create an array within this group. This method lightly wraps :func:`zarr.core.array.create_array`. @@ -1159,9 +1162,7 @@ async def create_array( ) @deprecated("Use AsyncGroup.create_array instead.", category=ZarrDeprecationWarning) - async def create_dataset( - self, name: str, *, shape: ShapeLike, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + async def create_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyAsyncArray: """Create an array. .. deprecated:: 3.0.0 @@ -1201,7 +1202,7 @@ async def require_dataset( dtype: npt.DTypeLike = None, exact: bool = False, **kwargs: Any, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Obtain an array, creating if it doesn't exist. .. deprecated:: 3.0.0 @@ -1238,7 +1239,7 @@ async def require_array( dtype: npt.DTypeLike = None, exact: bool = False, **kwargs: Any, - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Obtain an array, creating if it doesn't exist. Other `kwargs` are as per :func:`zarr.AsyncGroup.create_dataset`. @@ -1348,7 +1349,7 @@ async def members( *, use_consolidated_for_children: bool = True, ) -> AsyncGenerator[ - tuple[str, AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup], + tuple[str, AnyAsyncArray | AsyncGroup], None, ]: """ @@ -1387,7 +1388,7 @@ async def members( def _members_consolidated( self, max_depth: int | None, prefix: str = "" ) -> Generator[ - tuple[str, AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup], + tuple[str, AnyAsyncArray | AsyncGroup], None, ]: consolidated_metadata = self.metadata.consolidated_metadata @@ -1412,9 +1413,7 @@ def _members_consolidated( async def _members( self, max_depth: int | None, *, use_consolidated_for_children: bool = True - ) -> AsyncGenerator[ - tuple[str, AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup], None - ]: + ) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: skip_keys: tuple[str, ...] if self.metadata.zarr_format == 2: skip_keys = (".zattrs", ".zgroup", ".zarray", ".zmetadata") @@ -1454,9 +1453,7 @@ async def create_hierarchy( nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], *, overwrite: bool = False, - ) -> AsyncIterator[ - tuple[str, AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]] - ]: + ) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: """ Create a hierarchy of arrays or groups rooted at this group. @@ -1570,9 +1567,7 @@ async def group_values(self) -> AsyncGenerator[AsyncGroup, None]: async def arrays( self, - ) -> AsyncGenerator[ - tuple[str, AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]], None - ]: + ) -> AsyncGenerator[tuple[str, AnyAsyncArray], None]: """Iterate over arrays.""" async for key, value in self.members(): if isinstance(value, AsyncArray): @@ -1585,7 +1580,7 @@ async def array_keys(self) -> AsyncGenerator[str, None]: async def array_values( self, - ) -> AsyncGenerator[AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata], None]: + ) -> AsyncGenerator[AnyAsyncArray, None]: """Iterate over array values.""" async for _, array in self.arrays(): yield array @@ -1615,9 +1610,7 @@ async def tree(self, expand: bool | None = None, level: int | None = None) -> An raise NotImplementedError("'expand' is not yet implemented.") return await group_tree_async(self, max_depth=level) - async def empty( - self, *, name: str, shape: ChunkCoords, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + async def empty(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """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. @@ -1638,9 +1631,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 - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + async def zeros(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """Create an array, with zero being used as the default value for uninitialized portions of the array. Parameters @@ -1659,9 +1650,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 - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + async def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray: """Create an array, with one being used as the default value for uninitialized portions of the array. Parameters @@ -1682,7 +1671,7 @@ async def ones( async def full( self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. Parameters @@ -1711,7 +1700,7 @@ async def full( async def empty_like( self, *, name: str, data: async_api.ArrayLike, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create an empty sub-array like `data`. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -1733,7 +1722,7 @@ async def empty_like( async def zeros_like( self, *, name: str, data: async_api.ArrayLike, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create a sub-array of zeros like `data`. Parameters @@ -1754,7 +1743,7 @@ async def zeros_like( async def ones_like( self, *, name: str, data: async_api.ArrayLike, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create a sub-array of ones like `data`. Parameters @@ -1775,7 +1764,7 @@ async def ones_like( async def full_like( self, *, name: str, data: async_api.ArrayLike, **kwargs: Any - ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: + ) -> AnyAsyncArray: """Create a sub-array like `data` filled with the `fill_value` of `data` . Parameters @@ -1878,7 +1867,7 @@ def open( obj = sync(AsyncGroup.open(store, zarr_format=zarr_format)) return cls(obj) - def __getitem__(self, path: str) -> Array | Group: + def __getitem__(self, path: str) -> AnyArray | Group: """Obtain a group member. Parameters @@ -1911,7 +1900,7 @@ def __getitem__(self, path: str) -> Array | Group: else: return Group(obj) - def get(self, path: str, default: DefaultT | None = None) -> Array | Group | DefaultT | None: + def get(self, path: str, default: DefaultT | None = None) -> AnyArray | Group | DefaultT | None: """Obtain a group member, returning default if not found. Parameters @@ -2147,7 +2136,7 @@ def nmembers(self, max_depth: int | None = 0) -> int: def members( self, max_depth: int | None = 0, *, use_consolidated_for_children: bool = True - ) -> tuple[tuple[str, Array | Group], ...]: + ) -> tuple[tuple[str, AnyArray | Group], ...]: """ Returns an AsyncGenerator over the arrays and groups contained in this group. This method requires that `store_path.store` supports directory listing. @@ -2183,7 +2172,7 @@ def create_hierarchy( nodes: dict[str, ArrayV2Metadata | ArrayV3Metadata | GroupMetadata], *, overwrite: bool = False, - ) -> Iterator[tuple[str, Group | Array]]: + ) -> Iterator[tuple[str, Group | AnyArray]]: """ Create a hierarchy of arrays or groups rooted at this group. @@ -2322,7 +2311,7 @@ def group_values(self) -> Generator[Group, None]: for _, group in self.groups(): yield group - def arrays(self) -> Generator[tuple[str, Array], None]: + def arrays(self) -> Generator[tuple[str, AnyArray], None]: """Return the sub-arrays of this group as a generator of (name, array) pairs Examples @@ -2353,7 +2342,7 @@ def array_keys(self) -> Generator[str, None]: for name, _ in self.arrays(): yield name - def array_values(self) -> Generator[Array, None]: + def array_values(self) -> Generator[AnyArray, None]: """Return an iterator over group members. Examples @@ -2439,7 +2428,7 @@ def require_groups(self, *names: str) -> tuple[Group, ...]: """ return tuple(map(Group, self._sync(self._async_group.require_groups(*names)))) - def create(self, *args: Any, **kwargs: Any) -> Array: + def create(self, *args: Any, **kwargs: Any) -> AnyArray: # Backwards compatibility for 2.x return self.create_array(*args, **kwargs) @@ -2465,7 +2454,7 @@ def create_array( overwrite: bool = False, config: ArrayConfigLike | None = None, write_data: bool = True, - ) -> Array: + ) -> AnyArray: """Create an array within this group. This method lightly wraps :func:`zarr.core.array.create_array`. @@ -2594,7 +2583,7 @@ def create_array( ) @deprecated("Use Group.create_array instead.", category=ZarrDeprecationWarning) - def create_dataset(self, name: str, **kwargs: Any) -> Array: + def create_dataset(self, name: str, **kwargs: Any) -> AnyArray: """Create an array. .. deprecated:: 3.0.0 @@ -2618,7 +2607,7 @@ def create_dataset(self, name: str, **kwargs: Any) -> Array: return Array(self._sync(self._async_group.create_dataset(name, **kwargs))) @deprecated("Use Group.require_array instead.", category=ZarrDeprecationWarning) - def require_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> Array: + def require_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray: """Obtain an array, creating if it doesn't exist. .. deprecated:: 3.0.0 @@ -2642,7 +2631,7 @@ def require_dataset(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> Arra """ return Array(self._sync(self._async_group.require_array(name, shape=shape, **kwargs))) - def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> Array: + def require_array(self, name: str, *, shape: ShapeLike, **kwargs: Any) -> AnyArray: """Obtain an array, creating if it doesn't exist. Other `kwargs` are as per :func:`zarr.Group.create_array`. @@ -2660,7 +2649,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: ChunkCoords, **kwargs: Any) -> AnyArray: """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 +2670,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: ChunkCoords, **kwargs: Any) -> AnyArray: """Create an array, with zero being used as the default value for uninitialized portions of the array. Parameters @@ -2700,7 +2689,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: ChunkCoords, **kwargs: Any) -> AnyArray: """Create an array, with one being used as the default value for uninitialized portions of the array. Parameters @@ -2721,7 +2710,7 @@ def ones(self, *, name: str, shape: ChunkCoords, **kwargs: Any) -> Array: def full( self, *, name: str, shape: ChunkCoords, fill_value: Any | None, **kwargs: Any - ) -> Array: + ) -> AnyArray: """Create an array, with "fill_value" being used as the default value for uninitialized portions of the array. Parameters @@ -2746,7 +2735,7 @@ def full( ) ) - def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> Array: + def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: """Create an empty sub-array like `data`. The contents will be filled with the array's fill value or zeros if no fill value is provided. @@ -2772,7 +2761,7 @@ def empty_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> """ return Array(self._sync(self._async_group.empty_like(name=name, data=data, **kwargs))) - def zeros_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> Array: + def zeros_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: """Create a sub-array of zeros like `data`. Parameters @@ -2792,7 +2781,7 @@ def zeros_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> return Array(self._sync(self._async_group.zeros_like(name=name, data=data, **kwargs))) - def ones_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> Array: + def ones_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: """Create a sub-array of ones like `data`. Parameters @@ -2811,7 +2800,7 @@ def ones_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> A """ return Array(self._sync(self._async_group.ones_like(name=name, data=data, **kwargs))) - def full_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> Array: + def full_like(self, *, name: str, data: async_api.ArrayLike, **kwargs: Any) -> AnyArray: """Create a sub-array like `data` filled with the `fill_value` of `data` . Parameters @@ -2861,7 +2850,7 @@ def array( overwrite: bool = False, config: ArrayConfig | ArrayConfigLike | None = None, data: npt.ArrayLike | None = None, - ) -> Array: + ) -> AnyArray: """Create an array within this group. .. deprecated:: 3.0.0 @@ -2989,9 +2978,7 @@ async def create_hierarchy( store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], overwrite: bool = False, -) -> AsyncIterator[ - tuple[str, AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]] -]: +) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: """ Create a complete zarr hierarchy from a collection of metadata objects. @@ -3153,9 +3140,7 @@ async def create_nodes( *, store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], -) -> AsyncIterator[ - tuple[str, AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]] -]: +) -> AsyncIterator[tuple[str, AsyncGroup | AnyAsyncArray]]: """Create a collection of arrays and / or groups concurrently. Note: no attempt is made to validate that these arrays and / or groups collectively form a @@ -3333,7 +3318,7 @@ def _ensure_consistent_zarr_format( async def _getitem_semaphore( node: AsyncGroup, key: str, semaphore: asyncio.Semaphore | None -) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup: +) -> AnyAsyncArray | AsyncGroup: """ Wrap Group.getitem with an optional semaphore. @@ -3353,9 +3338,7 @@ async def _iter_members( node: AsyncGroup, skip_keys: tuple[str, ...], semaphore: asyncio.Semaphore | None, -) -> AsyncGenerator[ - tuple[str, AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup], None -]: +) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: """ Iterate over the arrays and groups contained in a group. @@ -3370,7 +3353,7 @@ async def _iter_members( Yields ------ - tuple[str, AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup] + tuple[str, AnyAsyncArray | AsyncGroup] """ # retrieve keys from storage @@ -3409,9 +3392,7 @@ async def _iter_members_deep( skip_keys: tuple[str, ...], semaphore: asyncio.Semaphore | None = None, use_consolidated_for_children: bool = True, -) -> AsyncGenerator[ - tuple[str, AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup], None -]: +) -> AsyncGenerator[tuple[str, AnyAsyncArray | AsyncGroup], None]: """ Iterate over the arrays and groups contained in a group, and optionally the arrays and groups contained in those groups. @@ -3434,7 +3415,7 @@ async def _iter_members_deep( Yields ------ - tuple[str, AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata] | AsyncGroup] + tuple[str, AnyAsyncArray | AsyncGroup] """ to_recurse = {} @@ -3449,7 +3430,7 @@ async def _iter_members_deep( if ( is_group and not use_consolidated_for_children - and node.metadata.consolidated_metadata is not None # type: ignore [union-attr] + and node.metadata.consolidated_metadata is not None ): node = cast("AsyncGroup", node) # We've decided not to trust consolidated metadata at this point, because we're @@ -3577,15 +3558,11 @@ def _build_metadata_v2( @overload -def _build_node( - *, store: Store, path: str, metadata: ArrayV2Metadata -) -> AsyncArray[ArrayV2Metadata]: ... +def _build_node(*, store: Store, path: str, metadata: ArrayV2Metadata) -> AsyncArrayV2: ... @overload -def _build_node( - *, store: Store, path: str, metadata: ArrayV3Metadata -) -> AsyncArray[ArrayV3Metadata]: ... +def _build_node(*, store: Store, path: str, metadata: ArrayV3Metadata) -> AsyncArrayV3: ... @overload @@ -3594,7 +3571,7 @@ def _build_node(*, store: Store, path: str, metadata: GroupMetadata) -> AsyncGro def _build_node( *, store: Store, path: str, metadata: ArrayV3Metadata | ArrayV2Metadata | GroupMetadata -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup: +) -> AnyAsyncArray | AsyncGroup: """ Take a metadata object and return a node (AsyncArray or AsyncGroup). """ @@ -3608,7 +3585,7 @@ def _build_node( raise ValueError(f"Unexpected metadata type: {type(metadata)}") # pragma: no cover -async def _get_node_v2(store: Store, path: str) -> AsyncArray[ArrayV2Metadata] | AsyncGroup: +async def _get_node_v2(store: Store, path: str) -> AsyncArrayV2 | AsyncGroup: """ Read a Zarr v2 AsyncArray or AsyncGroup from a path in a Store. @@ -3627,7 +3604,7 @@ async def _get_node_v2(store: Store, path: str) -> AsyncArray[ArrayV2Metadata] | return _build_node(store=store, path=path, metadata=metadata) -async def _get_node_v3(store: Store, path: str) -> AsyncArray[ArrayV3Metadata] | AsyncGroup: +async def _get_node_v3(store: Store, path: str) -> AsyncArrayV3 | AsyncGroup: """ Read a Zarr v3 AsyncArray or AsyncGroup from a path in a Store. @@ -3646,9 +3623,7 @@ async def _get_node_v3(store: Store, path: str) -> AsyncArray[ArrayV3Metadata] | return _build_node(store=store, path=path, metadata=metadata) -async def get_node( - store: Store, path: str, zarr_format: ZarrFormat -) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup: +async def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyAsyncArray | AsyncGroup: """ Get an AsyncArray or AsyncGroup from a path in a Store. @@ -3726,7 +3701,7 @@ async def create_rooted_hierarchy( store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], overwrite: bool = False, -) -> AsyncGroup | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]: +) -> AsyncGroup | AnyAsyncArray: """ Create an ``AsyncGroup`` or ``AsyncArray`` from a store and a dict of metadata documents. This function ensures that its input contains a specification of a root node, diff --git a/src/zarr/core/indexing.py b/src/zarr/core/indexing.py index 15cf6f0f1a..9465eead2b 100644 --- a/src/zarr/core/indexing.py +++ b/src/zarr/core/indexing.py @@ -30,10 +30,11 @@ from zarr.core.metadata import T_ArrayMetadata if TYPE_CHECKING: - from zarr.core.array import Array, AsyncArray + from zarr.core.array import AsyncArray from zarr.core.buffer import NDArrayLikeOrScalar from zarr.core.chunk_grids import ChunkGrid from zarr.core.common import ChunkCoords + from zarr.types import AnyArray IntSequence = list[int] | npt.NDArray[np.intp] @@ -79,7 +80,7 @@ def err_too_many_indices(selection: Any, shape: ChunkCoords) -> None: raise IndexError(f"too many indices for array; expected {len(shape)}, got {len(selection)}") -def _zarr_array_to_int_or_bool_array(arr: Array) -> npt.NDArray[np.intp] | npt.NDArray[np.bool_]: +def _zarr_array_to_int_or_bool_array(arr: AnyArray) -> npt.NDArray[np.intp] | npt.NDArray[np.bool_]: if arr.dtype.kind in ("i", "b"): return np.asarray(arr) else: @@ -931,10 +932,10 @@ def __iter__(self) -> Iterator[ChunkProjection]: @dataclass(frozen=True) class OIndex: - array: Array + array: AnyArray # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool - def __getitem__(self, selection: OrthogonalSelection | Array) -> NDArrayLikeOrScalar: + def __getitem__(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: from zarr.core.array import Array # if input is a Zarr array, we materialize it now. @@ -961,7 +962,7 @@ def __setitem__(self, selection: OrthogonalSelection, value: npt.ArrayLike) -> N class AsyncOIndex(Generic[T_ArrayMetadata]): array: AsyncArray[T_ArrayMetadata] - async def getitem(self, selection: OrthogonalSelection | Array) -> NDArrayLikeOrScalar: + async def getitem(self, selection: OrthogonalSelection | AnyArray) -> NDArrayLikeOrScalar: from zarr.core.array import Array # if input is a Zarr array, we materialize it now. @@ -1060,7 +1061,7 @@ def __iter__(self) -> Iterator[ChunkProjection]: @dataclass(frozen=True) class BlockIndex: - array: Array + array: AnyArray def __getitem__(self, selection: BasicSelection) -> NDArrayLikeOrScalar: fields, new_selection = pop_fields(selection) @@ -1249,11 +1250,11 @@ def __init__(self, selection: MaskSelection, shape: ChunkCoords, chunk_grid: Chu @dataclass(frozen=True) class VIndex: - array: Array + array: AnyArray # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool def __getitem__( - self, selection: CoordinateSelection | MaskSelection | Array + self, selection: CoordinateSelection | MaskSelection | AnyArray ) -> NDArrayLikeOrScalar: from zarr.core.array import Array @@ -1290,7 +1291,7 @@ class AsyncVIndex(Generic[T_ArrayMetadata]): # TODO: develop Array generic and move zarr.Array[np.intp] | zarr.Array[np.bool_] to ArrayOfIntOrBool async def getitem( - self, selection: CoordinateSelection | MaskSelection | Array + self, selection: CoordinateSelection | MaskSelection | AnyArray ) -> NDArrayLikeOrScalar: # TODO deduplicate these internals with the sync version of getitem # TODO requires solving this circular sync issue: https://github.com/zarr-developers/zarr-python/pull/3083#discussion_r2230737448 diff --git a/src/zarr/core/metadata/__init__.py b/src/zarr/core/metadata/__init__.py index 43b5ec98fe..87a66d23f3 100644 --- a/src/zarr/core/metadata/__init__.py +++ b/src/zarr/core/metadata/__init__.py @@ -5,7 +5,7 @@ ArrayMetadata: TypeAlias = ArrayV2Metadata | ArrayV3Metadata ArrayMetadataDict: TypeAlias = ArrayV2MetadataDict | ArrayV3MetadataDict -T_ArrayMetadata = TypeVar("T_ArrayMetadata", ArrayV2Metadata, ArrayV3Metadata) +T_ArrayMetadata = TypeVar("T_ArrayMetadata", ArrayV2Metadata, ArrayV3Metadata, covariant=True) __all__ = [ "ArrayMetadata", diff --git a/src/zarr/core/sync_group.py b/src/zarr/core/sync_group.py index 39d8a17992..417362b122 100644 --- a/src/zarr/core/sync_group.py +++ b/src/zarr/core/sync_group.py @@ -13,14 +13,14 @@ from collections.abc import Iterator from zarr.abc.store import Store - from zarr.core.array import Array from zarr.core.common import ZarrFormat from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata + from zarr.types import AnyArray def create_nodes( *, store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata] -) -> Iterator[tuple[str, Group | Array]]: +) -> Iterator[tuple[str, Group | AnyArray]]: """Create a collection of arrays and / or groups concurrently. Note: no attempt is made to validate that these arrays and / or groups collectively form a @@ -53,7 +53,7 @@ def create_hierarchy( store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], overwrite: bool = False, -) -> Iterator[tuple[str, Group | Array]]: +) -> Iterator[tuple[str, Group | AnyArray]]: """ Create a complete zarr hierarchy from a collection of metadata objects. @@ -115,7 +115,7 @@ def create_rooted_hierarchy( store: Store, nodes: dict[str, GroupMetadata | ArrayV2Metadata | ArrayV3Metadata], overwrite: bool = False, -) -> Group | Array: +) -> Group | AnyArray: """ Create a Zarr hierarchy with a root, and return the root node, which could be a ``Group`` or ``Array`` instance. @@ -140,7 +140,7 @@ def create_rooted_hierarchy( return _parse_async_node(async_node) -def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> Array | Group: +def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyArray | Group: """ Get an Array or Group from a path in a Store. diff --git a/src/zarr/testing/strategies.py b/src/zarr/testing/strategies.py index d0726c3dd9..862b38a562 100644 --- a/src/zarr/testing/strategies.py +++ b/src/zarr/testing/strategies.py @@ -23,6 +23,7 @@ from zarr.storage import MemoryStore, StoreLike from zarr.storage._common import _dereference_path from zarr.storage._utils import normalize_path +from zarr.types import AnyArray # Copied from Xarray _attr_keys = st.text(st.characters(), min_size=1) @@ -246,7 +247,7 @@ def arrays( arrays: st.SearchStrategy | None = None, attrs: st.SearchStrategy = attrs, zarr_formats: st.SearchStrategy = zarr_formats, -) -> Array: +) -> AnyArray: store = draw(stores, label="store") path = draw(paths, label="array parent") name = draw(array_names, label="array name") diff --git a/src/zarr/types.py b/src/zarr/types.py new file mode 100644 index 0000000000..38990982f9 --- /dev/null +++ b/src/zarr/types.py @@ -0,0 +1,23 @@ +from typing import Any, TypeAlias + +from zarr.core.array import Array, AsyncArray +from zarr.core.metadata.v2 import ArrayV2Metadata +from zarr.core.metadata.v3 import ArrayV3Metadata + +AnyAsyncArray: TypeAlias = AsyncArray[Any] +"""A Zarr format 2 or 3 `AsyncArray`""" + +AsyncArrayV2: TypeAlias = AsyncArray[ArrayV2Metadata] +"""A Zarr format 2 `AsyncArray`""" + +AsyncArrayV3: TypeAlias = AsyncArray[ArrayV3Metadata] +"""A Zarr format 3 `AsyncArray`""" + +AnyArray: TypeAlias = Array[Any] +"""A Zarr format 2 or 3 `Array`""" + +ArrayV2: TypeAlias = Array[ArrayV2Metadata] +"""A Zarr format 2 `Array`""" + +ArrayV3: TypeAlias = Array[ArrayV3Metadata] +"""A Zarr format 3 `Array`""" diff --git a/tests/test_api.py b/tests/test_api.py index 12acf80589..c7747f679d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -16,6 +16,7 @@ from zarr.abc.store import Store from zarr.core.common import JSON, MemoryOrder, ZarrFormat + from zarr.types import AnyArray import contextlib from typing import Literal @@ -128,7 +129,7 @@ def test_write_empty_chunks_warns(write_empty_chunks: bool, zarr_format: ZarrFor def test_open_normalized_path( memory_store: MemoryStore, path: str, node_type: Literal["array", "group"] ) -> None: - node: Group | Array + node: Group | AnyArray if node_type == "group": node = group(store=memory_store, path=path) elif node_type == "array": @@ -1385,7 +1386,7 @@ def test_no_overwrite_load(tmp_path: Path) -> None: zarr.zeros_like, ], ) -def test_auto_chunks(f: Callable[..., Array]) -> None: +def test_auto_chunks(f: Callable[..., AnyArray]) -> None: # Make sure chunks are set automatically across the public API # TODO: test shards with this test too shape = (1000, 1000) diff --git a/tests/test_array.py b/tests/test_array.py index 74201a4017..9d0f2ce1d7 100644 --- a/tests/test_array.py +++ b/tests/test_array.py @@ -7,7 +7,7 @@ import re import sys from itertools import accumulate -from typing import TYPE_CHECKING, Any, Literal +from typing import Any, Literal from unittest import mock import numcodecs @@ -19,7 +19,7 @@ import zarr.api.asynchronous import zarr.api.synchronous as sync_api from tests.conftest import skip_object_dtype -from zarr import Array, AsyncArray, Group +from zarr import Array, Group from zarr.abc.store import Store from zarr.codecs import ( BytesCodec, @@ -61,7 +61,6 @@ from zarr.core.group import AsyncGroup from zarr.core.indexing import BasicIndexer from zarr.core.metadata.v2 import ArrayV2Metadata -from zarr.core.metadata.v3 import ArrayV3Metadata from zarr.core.sync import sync from zarr.errors import ( ContainsArrayError, @@ -69,12 +68,10 @@ ZarrUserWarning, ) from zarr.storage import LocalStore, MemoryStore, StorePath +from zarr.types import AnyArray, AnyAsyncArray from .test_dtype.conftest import zdtype_examples -if TYPE_CHECKING: - from zarr.core.metadata.v3 import ArrayV3Metadata - @pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"]) @pytest.mark.parametrize("zarr_format", [2, 3]) @@ -356,9 +353,9 @@ def test_storage_transformers(store: MemoryStore, zarr_format: ZarrFormat | str) Array.from_dict(StorePath(store), data=metadata_dict) -@pytest.mark.parametrize("test_cls", [Array, AsyncArray[Any]]) +@pytest.mark.parametrize("test_cls", [AnyArray, AnyAsyncArray]) @pytest.mark.parametrize("nchunks", [2, 5, 10]) -def test_nchunks(test_cls: type[Array] | type[AsyncArray[Any]], nchunks: int) -> None: +def test_nchunks(test_cls: type[AnyArray] | type[AnyAsyncArray], nchunks: int) -> None: """ Test that nchunks returns the number of chunks defined for the array. """ @@ -373,8 +370,8 @@ def test_nchunks(test_cls: type[Array] | type[AsyncArray[Any]], nchunks: int) -> assert observed == expected -@pytest.mark.parametrize("test_cls", [Array, AsyncArray[Any]]) -async def test_nchunks_initialized(test_cls: type[Array] | type[AsyncArray[Any]]) -> None: +@pytest.mark.parametrize("test_cls", [AnyArray, AnyAsyncArray]) +async def test_nchunks_initialized(test_cls: type[AnyArray] | type[AnyAsyncArray]) -> None: """ Test that nchunks_initialized accurately returns the number of stored chunks. """ @@ -1347,7 +1344,7 @@ async def test_v2_chunk_encoding( # Normalize for property getters compressor_expected = () if compressor_expected is None else (compressor_expected,) - filters_expected = () if filters_expected is None else filters_expected + filters_expected = () if filters_expected is None else filters_expected # type: ignore[redundant-expr] assert arr.compressors == compressor_expected assert arr.filters == filters_expected @@ -1426,7 +1423,7 @@ async def test_with_data(impl: Literal["sync", "async"], store: Store) -> None: """ data = np.arange(10) name = "foo" - arr: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array + arr: AnyAsyncArray | AnyArray if impl == "sync": arr = sync_api.create_array(store, name=name, data=data) stored = arr[:] @@ -1647,7 +1644,7 @@ async def test_from_array_arraylike( store: Store, chunks: Literal["auto", "keep"] | tuple[int, int], write_data: bool, - src: Array | npt.ArrayLike, + src: AnyArray | npt.ArrayLike, ) -> None: fill_value = 42 result = zarr.from_array( @@ -1732,7 +1729,7 @@ def test_roundtrip_numcodecs() -> None: assert metadata["codecs"] == expected -def _index_array(arr: Array, index: Any) -> Any: +def _index_array(arr: AnyArray, index: Any) -> Any: return arr[index] diff --git a/tests/test_attributes.py b/tests/test_attributes.py index 4ce40e2cb0..269704d2a0 100644 --- a/tests/test_attributes.py +++ b/tests/test_attributes.py @@ -1,5 +1,5 @@ import json -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import pytest @@ -10,6 +10,9 @@ from tests.conftest import deep_nan_equal from zarr.core.common import ZarrFormat +if TYPE_CHECKING: + from zarr.types import AnyArray + @pytest.mark.parametrize("zarr_format", [2, 3]) @pytest.mark.parametrize( @@ -74,7 +77,7 @@ def test_update_no_changes() -> None: @pytest.mark.parametrize("group", [True, False]) def test_del_works(group: bool) -> None: store = zarr.storage.MemoryStore() - z: zarr.Group | zarr.Array + z: zarr.Group | AnyArray if group: z = zarr.create_group(store) else: @@ -84,7 +87,7 @@ def test_del_works(group: bool) -> None: del z.attrs["a"] assert dict(z.attrs) == {"c": 4} - z2: zarr.Group | zarr.Array + z2: zarr.Group | AnyArray if group: z2 = zarr.open_group(store) else: diff --git a/tests/test_codecs/test_codecs.py b/tests/test_codecs/test_codecs.py index a2dad41a1b..9e26ab2de0 100644 --- a/tests/test_codecs/test_codecs.py +++ b/tests/test_codecs/test_codecs.py @@ -29,11 +29,12 @@ from zarr.abc.store import Store from zarr.core.buffer.core import NDArrayLikeOrScalar from zarr.core.common import ChunkCoords, MemoryOrder + from zarr.types import AnyAsyncArray @dataclass(frozen=True) class _AsyncArrayProxy: - array: AsyncArray[Any] + array: AnyAsyncArray def __getitem__(self, selection: BasicSelection) -> _AsyncArraySelectionProxy: return _AsyncArraySelectionProxy(self.array, selection) @@ -41,7 +42,7 @@ def __getitem__(self, selection: BasicSelection) -> _AsyncArraySelectionProxy: @dataclass(frozen=True) class _AsyncArraySelectionProxy: - array: AsyncArray[Any] + array: AnyAsyncArray selection: BasicSelection async def get(self) -> NDArrayLikeOrScalar: diff --git a/tests/test_regression/test_v2_dtype_regression.py b/tests/test_regression/test_v2_dtype_regression.py index 9702ca7d23..85bf79cd19 100644 --- a/tests/test_regression/test_v2_dtype_regression.py +++ b/tests/test_regression/test_v2_dtype_regression.py @@ -13,11 +13,11 @@ import zarr.abc import zarr.abc.codec import zarr.codecs as zarrcodecs -from zarr.core.array import Array from zarr.core.chunk_key_encodings import V2ChunkKeyEncoding from zarr.core.dtype.npy.bytes import VariableLengthBytes from zarr.core.dtype.npy.string import VariableLengthUTF8 from zarr.storage import LocalStore +from zarr.types import ArrayV2, ArrayV3 if TYPE_CHECKING: from zarr.core.dtype import ZDTypeLike @@ -106,7 +106,7 @@ class ArrayParams: @pytest.fixture -def source_array_v2(tmp_path: Path, request: pytest.FixtureRequest) -> Array: +def source_array_v2(tmp_path: Path, request: pytest.FixtureRequest) -> ArrayV2: """ Writes a zarr array to a temporary directory based on the provided ArrayParams. The array is returned. @@ -144,7 +144,7 @@ def source_array_v2(tmp_path: Path, request: pytest.FixtureRequest) -> Array: @pytest.fixture -def source_array_v3(tmp_path: Path, request: pytest.FixtureRequest) -> Array: +def source_array_v3(tmp_path: Path, request: pytest.FixtureRequest) -> ArrayV3: """ Writes a zarr array to a temporary directory based on the provided ArrayParams. The array is returned. @@ -198,7 +198,7 @@ def source_array_v3(tmp_path: Path, request: pytest.FixtureRequest) -> Array: "source_array_v2", array_cases_v2_18, indirect=True, ids=tuple(map(str, array_cases_v2_18)) ) @pytest.mark.parametrize("script_path", script_paths) -def test_roundtrip_v2(source_array_v2: Array, tmp_path: Path, script_path: Path) -> None: +def test_roundtrip_v2(source_array_v2: ArrayV2, tmp_path: Path, script_path: Path) -> None: out_path = tmp_path / "out" copy_op = subprocess.run( [ @@ -222,7 +222,7 @@ def test_roundtrip_v2(source_array_v2: Array, tmp_path: Path, script_path: Path) @pytest.mark.parametrize( "source_array_v3", array_cases_v3_08, indirect=True, ids=tuple(map(str, array_cases_v3_08)) ) -def test_roundtrip_v3(source_array_v3: Array, tmp_path: Path) -> None: +def test_roundtrip_v3(source_array_v3: ArrayV3, tmp_path: Path) -> None: script_path = Path(__file__).resolve().parent / "scripts" / "v3.0.8.py" out_path = tmp_path / "out" copy_op = subprocess.run(