Add async array v2 and v3 types

Author: dstansby

src/zarr/core/array.py

@@ -144,7 +144,7 @@
     from zarr.codecs.sharding import ShardingCodecIndexLocation
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
     from zarr.storage import StoreLike
-    from zarr.types import AnyArray, AnyAsyncArray
+    from zarr.types import AnyArray, AnyAsyncArray, AsyncArrayV2, AsyncArrayV3
 
 
 # Array and AsyncArray are defined in the base ``zarr`` namespace
@@ -311,15 +311,15 @@ class AsyncArray(Generic[T_ArrayMetadata]):
 
     @overload
     def __init__(
-        self: AsyncArray[ArrayV2Metadata],
+        self: AsyncArrayV2,
         metadata: ArrayV2Metadata | ArrayV2MetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
     ) -> None: ...
 
     @overload
     def __init__(
-        self: AsyncArray[ArrayV3Metadata],
+        self: AsyncArrayV3,
         metadata: ArrayV3Metadata | ArrayV3MetadataDict,
         store_path: StorePath,
         config: ArrayConfigLike | None = None,
@@ -365,7 +365,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
@@ -394,7 +394,7 @@ async def create(
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
         config: ArrayConfigLike | None = None,
-    ) -> AsyncArray[ArrayV3Metadata]: ...
+    ) -> AsyncArrayV3: ...
 
     @overload
     @classmethod
@@ -422,7 +422,7 @@ async def create(
         overwrite: bool = False,
         data: npt.ArrayLike | None = None,
         config: ArrayConfigLike | None = None,
-    ) -> AsyncArray[ArrayV3Metadata]: ...
+    ) -> AsyncArrayV3: ...
 
     @overload
     @classmethod
@@ -806,7 +806,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()
@@ -887,7 +887,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()
@@ -947,7 +947,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

src/zarr/core/group.py

@@ -79,7 +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
+    from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3
 
 logger = logging.getLogger("zarr.group")
 
@@ -114,11 +114,11 @@ def parse_attributes(data: Any) -> dict[str, Any]:
 
 
 @overload
-def _parse_async_node(node: AsyncArray[ArrayV3Metadata]) -> ArrayV3: ...
+def _parse_async_node(node: AsyncArrayV3) -> ArrayV3: ...
 
 
 @overload
-def _parse_async_node(node: AsyncArray[ArrayV2Metadata]) -> ArrayV2: ...
+def _parse_async_node(node: AsyncArrayV2) -> ArrayV2: ...
 
 
 @overload
@@ -3706,15 +3706,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
@@ -3737,7 +3733,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.
 
@@ -3756,7 +3752,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.
 

src/zarr/types.py

@@ -7,6 +7,11 @@
 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`"""