Parametrize Array with v2/v3 metadata

Author: dstansby

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",
@@ -168,7 +169,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
@@ -365,7 +366,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
@@ -633,7 +634,7 @@ def create(
     storage_options: dict[str, Any] | None = None,
     config: ArrayConfigLike | None = None,
     **kwargs: Any,
-) -> Array:
+) -> AnyArray:
     """Create an array.
 
     Parameters
@@ -769,7 +770,7 @@ def create_array(
     overwrite: bool = False,
     config: ArrayConfigLike | None = None,
     write_data: bool = True,
-) -> Array:
+) -> AnyArray:
     """Create an array.
 
     This function wraps :func:`zarr.core.array.create_array`.
@@ -934,7 +935,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
@@ -1128,7 +1129,7 @@ def from_array(
 
 
 # TODO: add type annotations for kwargs
-def empty(shape: ChunkCoords, **kwargs: Any) -> Array:
+def empty(shape: ChunkCoords, **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.
 
@@ -1155,7 +1156,7 @@ def empty(shape: ChunkCoords, **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.
 
@@ -1181,7 +1182,7 @@ def empty_like(a: ArrayLike, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs and fill_value
-def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> Array:
+def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AnyArray:
     """Create an array with a default fill value.
 
     Parameters
@@ -1203,7 +1204,7 @@ def full(shape: ChunkCoords, 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
@@ -1222,7 +1223,7 @@ def full_like(a: ArrayLike, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs
-def ones(shape: ChunkCoords, **kwargs: Any) -> Array:
+def ones(shape: ChunkCoords, **kwargs: Any) -> AnyArray:
     """Create an array with a fill value of one.
 
     Parameters
@@ -1241,7 +1242,7 @@ def ones(shape: ChunkCoords, **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
@@ -1267,7 +1268,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
@@ -1303,7 +1304,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
@@ -1324,7 +1325,7 @@ def open_like(a: ArrayLike, path: str, **kwargs: Any) -> Array:
 
 
 # TODO: add type annotations for kwargs
-def zeros(shape: ChunkCoords, **kwargs: Any) -> Array:
+def zeros(shape: ChunkCoords, **kwargs: Any) -> AnyArray:
     """Create an array with a fill value of zero.
 
     Parameters
@@ -1343,7 +1344,7 @@ def zeros(shape: ChunkCoords, **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

@@ -1782,12 +1782,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.")

src/zarr/types.py

@@ -0,0 +1,14 @@
+from typing import Any
+
+from zarr.core.array import Array
+from zarr.core.metadata.v2 import ArrayV2Metadata
+from zarr.core.metadata.v3 import ArrayV3Metadata
+
+AnyArray = Array[Any]
+AnyArray.__doc__ = "A Zarr format 2 or 3 `Array`"
+
+ArrayV2 = Array[ArrayV2Metadata]
+ArrayV2.__doc__ = "A Zarr format 2 `Array`"
+
+ArrayV3 = Array[ArrayV3Metadata]
+ArrayV3.__doc__ = "A Zarr format 3 `Array`"

tests/test_api.py

@@ -15,6 +15,7 @@
 
     from zarr.abc.store import Store
     from zarr.core.common import JSON, MemoryOrder, ZarrFormat
+    from zarr.types import AnyArray
 
 import contextlib
 from typing import Literal
@@ -127,7 +128,7 @@ def test_write_empty_chunks_warns(write_empty_chunks: bool, zarr_format: ZarrFor
 def test_open_normalized_path(
     memory_store: MemoryStore, path: str, node_type: Literal["array", "group"]
 ) -> None:
-    node: Group | Array
+    node: Group | AnyArray
     if node_type == "group":
         node = group(store=memory_store, path=path)
     elif node_type == "array":
@@ -1384,7 +1385,7 @@ def test_no_overwrite_load(tmp_path: Path) -> None:
         zarr.zeros_like,
     ],
 )
-def test_auto_chunks(f: Callable[..., Array]) -> None:
+def test_auto_chunks(f: Callable[..., AnyArray]) -> None:
     # Make sure chunks are set automatically across the public API
     # TODO: test shards with this test too
     shape = (1000, 1000)

tests/test_array.py

@@ -65,6 +65,7 @@
 from zarr.core.sync import sync
 from zarr.errors import ContainsArrayError, ContainsGroupError
 from zarr.storage import LocalStore, MemoryStore, StorePath
+from zarr.types import AnyArray
 
 from .test_dtype.conftest import zdtype_examples
 
@@ -1616,7 +1617,7 @@ async def test_from_array_arraylike(
     store: Store,
     chunks: Literal["auto", "keep"] | tuple[int, int],
     write_data: bool,
-    src: Array | npt.ArrayLike,
+    src: AnyArray | npt.ArrayLike,
 ) -> None:
     fill_value = 42
     result = zarr.from_array(
@@ -1698,7 +1699,7 @@ def test_roundtrip_numcodecs() -> None:
     assert metadata["codecs"] == expected
 
 
-def _index_array(arr: Array, index: Any) -> Any:
+def _index_array(arr: AnyArray, index: Any) -> Any:
     return arr[index]