Add an AnyAsyncArray type

Author: dstansby

src/zarr/api/asynchronous.py

@@ -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
 from zarr.storage import StorePath
 from zarr.storage._common import make_store_path
@@ -52,12 +52,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
+    from zarr.types import AnyArray, AnyAsyncArray
 
     # TODO: this type could use some more thought
-    ArrayLike: TypeAlias = (
-        AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AnyArray | npt.NDArray[Any]
-    )
+    ArrayLike: TypeAlias = AnyAsyncArray | AnyArray | npt.NDArray[Any]
     PathLike = str
 
 __all__ = [
@@ -309,7 +307,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
@@ -564,9 +562,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 | AnyArray, **kwargs: Any
-) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
+async def array(data: npt.ArrayLike | AnyArray, **kwargs: Any) -> AnyAsyncArray:
     """Create an array filled with `data`.
 
     Parameters
@@ -898,7 +894,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
@@ -1070,9 +1066,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.
 
@@ -1093,9 +1087,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.
 
@@ -1122,9 +1114,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.
 
@@ -1146,9 +1136,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
@@ -1169,9 +1157,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.
 
@@ -1190,9 +1176,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
@@ -1219,7 +1203,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
@@ -1267,9 +1251,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
@@ -1292,9 +1274,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.
 
@@ -1313,9 +1293,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

src/zarr/core/array.py

@@ -140,7 +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
+    from zarr.types import AnyArray, AnyAsyncArray
 
 
 # Array and AsyncArray are defined in the base ``zarr`` namespace
@@ -439,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.")
@@ -473,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
@@ -614,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`.
@@ -636,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(
@@ -905,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.
 
@@ -937,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.
 
@@ -3833,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.
@@ -3865,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
 
@@ -3947,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
@@ -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