Parametrize Array with v2/v3 metadata

Author: dstansby

src/zarr/api/asynchronous.py

@@ -3,7 +3,7 @@
 import asyncio
 import dataclasses
 import warnings
-from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypeAlias, TypedDict, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -58,9 +58,12 @@
     from zarr.core.chunk_key_encodings import ChunkKeyEncoding
     from zarr.core.metadata.v2 import CompressorLikev2
     from zarr.storage import StoreLike
+    from zarr.types import AnyArray
 
     # TODO: this type could use some more thought
-    ArrayLike = AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | Array | npt.NDArray[Any]
+    ArrayLike: TypeAlias = (
+        AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AnyArray | npt.NDArray[Any]
+    )
     PathLike = str
 
 __all__ = [
@@ -602,7 +605,7 @@ async def tree(grp: AsyncGroup, expand: bool | None = None, level: int | None =
 
 
 async def array(
-    data: npt.ArrayLike | Array, **kwargs: Any
+    data: npt.ArrayLike | AnyArray, **kwargs: Any
 ) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
     """Create an array filled with `data`.
 

src/zarr/api/synchronous.py

@@ -40,6 +40,7 @@
     )
     from zarr.core.dtype import ZDTypeLike
     from zarr.storage import StoreLike
+    from zarr.types import AnyArray
 
 __all__ = [
     "array",
@@ -181,7 +182,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
@@ -387,7 +388,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
@@ -652,7 +653,7 @@ def create(
     storage_options: dict[str, Any] | None = None,
     config: ArrayConfigLike | None = None,
     **kwargs: Any,
-) -> Array:
+) -> AnyArray:
     """Create an array.
 
     Parameters
@@ -836,7 +837,7 @@ def create_array(
     overwrite: bool = False,
     config: ArrayConfigLike | None = None,
     write_data: bool = True,
-) -> Array:
+) -> AnyArray:
     """Create an array.
 
     This function wraps [zarr.core.array.create_array][].
@@ -987,7 +988,7 @@ def create_array(
 def from_array(
     store: StoreLike,
     *,
-    data: Array | npt.ArrayLike,
+    data: AnyArray | npt.ArrayLike,
     write_data: bool = True,
     name: str | None = None,
     chunks: Literal["auto", "keep"] | tuple[int, ...] = "keep",
@@ -1004,7 +1005,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
@@ -1203,7 +1204,7 @@ def from_array(
 
 
 # TODO: add type annotations for kwargs
-def empty(shape: tuple[int, ...], **kwargs: Any) -> Array:
+def empty(shape: tuple[int, ...], **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.
 
@@ -1230,7 +1231,7 @@ def empty(shape: tuple[int, ...], **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.
 
@@ -1256,7 +1257,7 @@ def empty_like(a: ArrayLike, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs and fill_value
-def full(shape: tuple[int, ...], fill_value: Any, **kwargs: Any) -> Array:
+def full(shape: tuple[int, ...], fill_value: Any, **kwargs: Any) -> AnyArray:
     """Create an array with a default fill value.
 
     Parameters
@@ -1278,7 +1279,7 @@ def full(shape: tuple[int, ...], 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
@@ -1297,7 +1298,7 @@ def full_like(a: ArrayLike, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs
-def ones(shape: tuple[int, ...], **kwargs: Any) -> Array:
+def ones(shape: tuple[int, ...], **kwargs: Any) -> AnyArray:
     """Create an array with a fill value of one.
 
     Parameters
@@ -1316,7 +1317,7 @@ def ones(shape: tuple[int, ...], **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
@@ -1343,7 +1344,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
@@ -1385,7 +1386,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
@@ -1406,7 +1407,7 @@ def open_like(a: ArrayLike, path: str, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs
-def zeros(shape: tuple[int, ...], **kwargs: Any) -> Array:
+def zeros(shape: tuple[int, ...], **kwargs: Any) -> AnyArray:
     """Create an array with a fill value of zero.
 
     Parameters
@@ -1425,7 +1426,7 @@ def zeros(shape: tuple[int, ...], **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

src/zarr/core/array.py

@@ -144,6 +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
 
 
 # Array and AsyncArray are defined in the base ``zarr`` namespace
@@ -1989,12 +1990,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)
@@ -2027,7 +2028,7 @@ def create(
         # runtime
         overwrite: bool = False,
         config: ArrayConfigLike | None = None,
-    ) -> Array:
+    ) -> AnyArray:
         """Creates a new Array instance from an initialized store.
 
         !!! warning "Deprecated"
@@ -2171,7 +2172,7 @@ def _create(
         # runtime
         overwrite: bool = False,
         config: ArrayConfigLike | None = None,
-    ) -> Array:
+    ) -> AnyArray:
         """Creates a new Array instance from an initialized store.
         Deprecated in favor of [`zarr.create_array`][].
         """
@@ -2196,14 +2197,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.
 
@@ -2227,13 +2228,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
@@ -2249,7 +2250,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:
@@ -4085,7 +4086,7 @@ def append(self, data: npt.ArrayLike, axis: int = 0) -> tuple[int, ...]:
         """
         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.
 
@@ -4110,11 +4111,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"<Array {self.store_path} shape={self.shape} dtype={self.dtype}>"
@@ -4236,7 +4234,7 @@ class ShardsConfigParam(TypedDict):
 async def from_array(
     store: StoreLike,
     *,
-    data: Array | npt.ArrayLike,
+    data: AnyArray | npt.ArrayLike,
     write_data: bool = True,
     name: str | None = None,
     chunks: Literal["auto", "keep"] | tuple[int, ...] = "keep",
@@ -4478,7 +4476,7 @@ async def from_array(
         if isinstance(data, Array):
 
             async def _copy_array_region(
-                chunk_coords: tuple[int, ...] | slice, _data: Array
+                chunk_coords: tuple[int, ...] | slice, _data: AnyArray
             ) -> None:
                 arr = await _data._async_array.getitem(chunk_coords)
                 await result.setitem(chunk_coords, arr)
@@ -4908,7 +4906,7 @@ async def create_array(
 
 
 def _parse_keep_array_attr(
-    data: Array | npt.ArrayLike,
+    data: AnyArray | npt.ArrayLike,
     chunks: Literal["auto", "keep"] | tuple[int, ...],
     shards: ShardsLike | None | Literal["keep"],
     filters: FiltersLike | Literal["keep"],
@@ -5287,7 +5285,7 @@ def _parse_data_params(
 
 
 def _iter_chunk_coords(
-    array: Array | AsyncArray[Any],
+    array: AnyArray | AsyncArray[Any],
     *,
     origin: Sequence[int] | None = None,
     selection_shape: Sequence[int] | None = None,
@@ -5318,7 +5316,7 @@ def _iter_chunk_coords(
 
 
 def _iter_shard_coords(
-    array: Array | AsyncArray[Any],
+    array: AnyArray | AsyncArray[Any],
     *,
     origin: Sequence[int] | None = None,
     selection_shape: Sequence[int] | None = None,
@@ -5349,7 +5347,7 @@ def _iter_shard_coords(
 
 
 def _iter_shard_keys(
-    array: Array | AsyncArray[Any],
+    array: AnyArray | AsyncArray[Any],
     *,
     origin: Sequence[int] | None = None,
     selection_shape: Sequence[int] | None = None,
@@ -5378,7 +5376,7 @@ def _iter_shard_keys(
 
 
 def _iter_shard_regions(
-    array: Array | AsyncArray[Any],
+    array: AnyArray | AsyncArray[Any],
     *,
     origin: Sequence[int] | None = None,
     selection_shape: Sequence[int] | None = None,
@@ -5413,7 +5411,7 @@ def _iter_shard_regions(
 
 
 def _iter_chunk_regions(
-    array: Array | AsyncArray[Any],
+    array: AnyArray | AsyncArray[Any],
     *,
     origin: Sequence[int] | None = None,
     selection_shape: Sequence[int] | None = None,

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