make asyncarray generic w.r.t metadata

Author: d-v-b

src/zarr/api/asynchronous.py

@@ -28,7 +28,7 @@
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
 
     # TODO: this type could use some more thought
-    ArrayLike = AsyncArray | Array | npt.NDArray[Any]
+    ArrayLike = AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array | npt.NDArray[Any]
     PathLike = str
 
 __all__ = [
@@ -198,7 +198,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 | AsyncGroup:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup:
     """Convenience function to open a group or array using file-mode-like semantics.
 
     Parameters
@@ -400,7 +400,9 @@ async def tree(*args: Any, **kwargs: Any) -> None:
     raise NotImplementedError
 
 
-async def array(data: npt.ArrayLike, **kwargs: Any) -> AsyncArray:
+async def array(
+    data: npt.ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array filled with `data`.
 
     Parameters
@@ -652,7 +654,7 @@ async def create(
     dimension_names: Iterable[str] | None = None,
     storage_options: dict[str, Any] | None = None,
     **kwargs: Any,
-) -> AsyncArray:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array.
 
     Parameters
@@ -804,7 +806,9 @@ async def create(
     )
 
 
-async def empty(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def empty(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array.
 
     Parameters
@@ -823,7 +827,9 @@ async def empty(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=None, **kwargs)
 
 
-async def empty_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def empty_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an empty array like `a`.
 
     Parameters
@@ -843,7 +849,9 @@ async def empty_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
 
 
 # TODO: add type annotations for fill_value and kwargs
-async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AsyncArray:
+async def full(
+    shape: ChunkCoords, 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.
 
@@ -865,7 +873,9 @@ async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AsyncArray
 
 
 # TODO: add type annotations for kwargs
-async def full_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def full_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create a filled array like `a`.
 
     Parameters
@@ -886,7 +896,9 @@ async def full_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
     return await full(**like_kwargs)
 
 
-async def ones(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def ones(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with one being used as the default value for
     uninitialized portions of the array.
 
@@ -905,7 +917,9 @@ async def ones(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=1, **kwargs)
 
 
-async def ones_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def ones_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array of ones like `a`.
 
     Parameters
@@ -932,7 +946,7 @@ async def open_array(
     path: PathLike | None = None,
     storage_options: dict[str, Any] | None = None,
     **kwargs: Any,  # TODO: type kwargs as valid args to save
-) -> AsyncArray:
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Open an array using file-mode-like semantics.
 
     Parameters
@@ -975,7 +989,9 @@ async def open_array(
         raise
 
 
-async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AsyncArray:
+async def open_like(
+    a: ArrayLike, path: str, **kwargs: Any
+) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
     """Open a persistent array like `a`.
 
     Parameters
@@ -998,7 +1014,9 @@ async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AsyncArray:
     return await open_array(path=path, **like_kwargs)
 
 
-async def zeros(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
+async def zeros(
+    shape: ChunkCoords, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array, with zero being used as the default value for
     uninitialized portions of the array.
 
@@ -1017,7 +1035,9 @@ async def zeros(shape: ChunkCoords, **kwargs: Any) -> AsyncArray:
     return await create(shape=shape, fill_value=0, **kwargs)
 
 
-async def zeros_like(a: ArrayLike, **kwargs: Any) -> AsyncArray:
+async def zeros_like(
+    a: ArrayLike, **kwargs: Any
+) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array of zeros like `a`.
 
     Parameters

src/zarr/core/array.py

@@ -4,7 +4,7 @@
 from asyncio import gather
 from dataclasses import dataclass, field, replace
 from logging import getLogger
-from typing import TYPE_CHECKING, Any, Literal, cast
+from typing import TYPE_CHECKING, Any, Generic, Literal, TypeVar, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -73,10 +73,10 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable, Iterator, Sequence
+    from typing import Self
 
     from zarr.abc.codec import Codec, CodecPipeline
     from zarr.core.group import AsyncGroup
-    from zarr.core.metadata.common import ArrayMetadata
 
 # Array and AsyncArray are defined in the base ``zarr`` namespace
 __all__ = ["create_codec_pipeline", "parse_array_metadata"]
@@ -160,16 +160,19 @@ async def get_array_metadata(
     return metadata_dict
 
 
+TArrayMeta = TypeVar("TArrayMeta", ArrayV2Metadata, ArrayV3Metadata)
+
+
 @dataclass(frozen=True)
-class AsyncArray:
-    metadata: ArrayMetadata
+class AsyncArray(Generic[TArrayMeta]):
+    metadata: TArrayMeta
     store_path: StorePath
     codec_pipeline: CodecPipeline = field(init=False)
     order: Literal["C", "F"]
 
     def __init__(
         self,
-        metadata: ArrayMetadata | dict[str, Any],
+        metadata: ArrayV2Metadata | ArrayV3Metadata | dict[str, Any],
         store_path: StorePath,
         order: Literal["C", "F"] | None = None,
     ) -> None:
@@ -218,7 +221,7 @@ async def create(
         # runtime
         exists_ok: bool = False,
         data: npt.ArrayLike | None = None,
-    ) -> AsyncArray:
+    ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
         store_path = await make_store_path(store)
 
         shape = parse_shapelike(shape)
@@ -231,7 +234,7 @@ async def create(
             _chunks = normalize_chunks(chunks, shape, dtype.itemsize)
         else:
             _chunks = normalize_chunks(chunk_shape, shape, dtype.itemsize)
-
+        result: AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]
         if zarr_format == 3:
             if dimension_separator is not None:
                 raise ValueError(
@@ -313,7 +316,7 @@ async def _create_v3(
         dimension_names: Iterable[str] | None = None,
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
-    ) -> AsyncArray:
+    ) -> AsyncArray[ArrayV3Metadata]:
         if not exists_ok:
             await ensure_no_existing_node(store_path, zarr_format=3)
 
@@ -344,7 +347,9 @@ async def _create_v3(
 
         array = cls(metadata=metadata, store_path=store_path)
         await array._save_metadata(metadata, ensure_parents=True)
-        return array
+        # type inference is inconsistent here and seems to conclude
+        # that array has type Array[ArrayV2Metadata]
+        return array  # type: ignore[return-value]
 
     @classmethod
     async def _create_v2(
@@ -361,7 +366,7 @@ async def _create_v2(
         compressor: dict[str, JSON] | None = None,
         attributes: dict[str, JSON] | None = None,
         exists_ok: bool = False,
-    ) -> AsyncArray:
+    ) -> AsyncArray[ArrayV2Metadata]:
         if not exists_ok:
             await ensure_no_existing_node(store_path, zarr_format=2)
         if order is None:
@@ -383,14 +388,14 @@ async def _create_v2(
         )
         array = cls(metadata=metadata, store_path=store_path)
         await array._save_metadata(metadata, ensure_parents=True)
-        return array
+        return array  # type: ignore[return-value]
 
     @classmethod
     def from_dict(
         cls,
         store_path: StorePath,
         data: dict[str, JSON],
-    ) -> AsyncArray:
+    ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
         metadata = parse_array_metadata(data)
         return cls(metadata=metadata, store_path=store_path)
 
@@ -399,7 +404,7 @@ async def open(
         cls,
         store: StoreLike,
         zarr_format: ZarrFormat | None = 3,
-    ) -> AsyncArray:
+    ) -> AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]:
         store_path = await make_store_path(store)
         metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format)
         return cls(store_path=store_path, metadata=metadata_dict)
@@ -631,7 +636,9 @@ async def getitem(
         )
         return await self._get_selection(indexer, prototype=prototype)
 
-    async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = False) -> None:
+    async def _save_metadata(
+        self, metadata: ArrayV2Metadata | ArrayV3Metadata, ensure_parents: bool = False
+    ) -> None:
         to_save = metadata.to_buffer_dict(default_buffer_prototype())
         awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
 
@@ -719,9 +726,7 @@ async def setitem(
         )
         return await self._set_selection(indexer, value, prototype=prototype)
 
-    async def resize(
-        self, new_shape: ChunkCoords, delete_outside_chunks: bool = True
-    ) -> AsyncArray:
+    async def resize(self, new_shape: ChunkCoords, delete_outside_chunks: bool = True) -> Self:
         assert len(new_shape) == len(self.metadata.shape)
         new_metadata = self.metadata.update_shape(new_shape)
 
@@ -747,7 +752,7 @@ async def _delete_key(key: str) -> None:
         await self._save_metadata(new_metadata)
         return replace(self, metadata=new_metadata)
 
-    async def update_attributes(self, new_attributes: dict[str, JSON]) -> AsyncArray:
+    async def update_attributes(self, new_attributes: dict[str, JSON]) -> Self:
         new_metadata = self.metadata.update_attributes(new_attributes)
 
         # Write new metadata
@@ -763,7 +768,7 @@ async def info(self) -> None:
 
 @dataclass(frozen=True)
 class Array:
-    _async_array: AsyncArray
+    _async_array: AsyncArray[ArrayV3Metadata] | AsyncArray[ArrayV2Metadata]
 
     @classmethod
     @_deprecate_positional_args
@@ -879,7 +884,7 @@ def basename(self) -> str | None:
         return self._async_array.basename
 
     @property
-    def metadata(self) -> ArrayMetadata:
+    def metadata(self) -> ArrayV2Metadata | ArrayV3Metadata:
         return self._async_array.metadata
 
     @property
@@ -2309,18 +2314,10 @@ def resize(self, new_shape: ChunkCoords) -> Array:
         the data falling outside the new array but inside the boundary chunks
         would be restored by a subsequent resize operation that grows the array size.
         """
-        return type(self)(
-            sync(
-                self._async_array.resize(new_shape),
-            )
-        )
+        return type(self)(sync(self._async_array.resize(new_shape)))  # type: ignore[arg-type]
 
     def update_attributes(self, new_attributes: dict[str, JSON]) -> Array:
-        return type(self)(
-            sync(
-                self._async_array.update_attributes(new_attributes),
-            )
-        )
+        return type(self)(sync(self._async_array.update_attributes(new_attributes)))  # type: ignore[arg-type]
 
     def __repr__(self) -> str:
         return f"<Array {self.store_path} shape={self.shape} dtype={self.dtype}>"
@@ -2331,7 +2328,9 @@ def info(self) -> None:
         )
 
 
-def nchunks_initialized(array: AsyncArray | Array) -> int:
+def nchunks_initialized(
+    array: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array,
+) -> int:
     """
     Calculate the number of chunks that have been initialized, i.e. the number of chunks that have
     been persisted to the storage backend.
@@ -2353,7 +2352,9 @@ def nchunks_initialized(array: AsyncArray | Array) -> int:
     return len(chunks_initialized(array))
 
 
-def chunks_initialized(array: Array | AsyncArray) -> tuple[str, ...]:
+def chunks_initialized(
+    array: Array | AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata],
+) -> tuple[str, ...]:
     """
     Return the keys of the chunks that have been persisted to the storage backend.
 
@@ -2385,7 +2386,9 @@ def chunks_initialized(array: Array | AsyncArray) -> tuple[str, ...]:
     return tuple(out)
 
 
-def _build_parents(node: AsyncArray | AsyncGroup) -> list[AsyncGroup]:
+def _build_parents(
+    node: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup,
+) -> list[AsyncGroup]:
     from zarr.core.group import AsyncGroup, GroupMetadata
 
     store = node.store_path.store