Parametrize Array with v2/v3 metadata (#3304)

* Parametrize Array with v2/v3 metadata

* Try fixing docstrings

* Add an AnyAsyncArray type

* Add async array v2 and v3 types

* Add changelog entry

* Pass sub-classes through returns.

* Chage release note to MarkDown

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: dstansby

changes/3304.feature.md

@@ -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.

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
@@ -37,7 +37,7 @@
     GroupMetadata,
     create_hierarchy,
 )
-from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
+from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata
 from zarr.errors import (
     ArrayNotFoundError,
     GroupNotFoundError,
@@ -58,9 +58,10 @@
     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, 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__ = [
@@ -335,7 +336,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
@@ -601,9 +602,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
@@ -919,7 +918,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
@@ -1106,9 +1105,7 @@ async def create(
     )
 
 
-async def empty(
-    shape: tuple[int, ...], **kwargs: Any
-) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
+async def empty(shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray:
     """Create an empty array with the specified shape. The contents will be filled with the
     specified fill value or zeros if no fill value is provided.
 
@@ -1128,9 +1125,7 @@ async def empty(
     return await create(shape=shape, **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.
 
@@ -1159,9 +1154,7 @@ async def empty_like(
 
 
 # TODO: add type annotations for fill_value and kwargs
-async def full(
-    shape: tuple[int, ...], fill_value: Any, **kwargs: Any
-) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
+async def full(shape: tuple[int, ...], fill_value: Any, **kwargs: Any) -> AnyAsyncArray:
     """Create an array, with `fill_value` being used as the default value for
     uninitialized portions of the array.
 
@@ -1183,9 +1176,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
@@ -1206,9 +1197,7 @@ async def full_like(
     return await full(**like_kwargs)  # type: ignore[arg-type]
 
 
-async def ones(
-    shape: tuple[int, ...], **kwargs: Any
-) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
+async def ones(shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray:
     """Create an array, with one being used as the default value for
     uninitialized portions of the array.
 
@@ -1227,9 +1216,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
@@ -1256,7 +1243,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
@@ -1307,9 +1294,7 @@ async def open_array(
         raise ArrayNotFoundError(msg) from err
 
 
-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
@@ -1332,9 +1317,7 @@ async def open_like(
     return await open_array(path=path, **like_kwargs)  # type: ignore[arg-type]
 
 
-async def zeros(
-    shape: tuple[int, ...], **kwargs: Any
-) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
+async def zeros(shape: tuple[int, ...], **kwargs: Any) -> AnyAsyncArray:
     """Create an array, with zero being used as the default value for
     uninitialized portions of the array.
 
@@ -1353,9 +1336,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/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][].
@@ -989,7 +990,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",
@@ -1006,7 +1007,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
@@ -1220,7 +1221,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.
 
@@ -1247,7 +1248,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.
 
@@ -1273,7 +1274,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
@@ -1295,7 +1296,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
@@ -1314,7 +1315,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
@@ -1333,7 +1334,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
@@ -1360,7 +1361,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
@@ -1402,7 +1403,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
@@ -1423,7 +1424,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
@@ -1442,7 +1443,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