Skip to content

Parametrize Array with v2/v3 metadata #3304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
+208 −222
Open

Parametrize Array with v2/v3 metadata #3304

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
97d0dd7
Parametrize Array with v2/v3 metadata
dstansby Jul 28, 2025
b571f5d
Try fixing docstrings
dstansby Jul 28, 2025
2b418c0
Add an AnyAsyncArray type
dstansby Jul 28, 2025
1a46787
Add async array v2 and v3 types
dstansby Jul 28, 2025
2211f88
Add changelog entry
dstansby Jul 28, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions changes/3304.feature.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
53 changes: 17 additions & 36 deletions src/zarr/api/asynchronous.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import asyncio
import dataclasses
import warnings
from typing import TYPE_CHECKING, Any, Literal, cast
from typing import TYPE_CHECKING, Any, Literal, TypeAlias, cast

import numpy as np
import numpy.typing as npt
Expand Down Expand Up @@ -38,7 +38,7 @@
GroupMetadata,
create_hierarchy,
)
from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata, ArrayV3Metadata
from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata
from zarr.errors import (
GroupNotFoundError,
NodeTypeValidationError,
Expand All @@ -58,9 +58,10 @@
from zarr.core.buffer import NDArrayLikeOrScalar
from zarr.core.chunk_key_encodings import ChunkKeyEncoding
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__ = [
Expand Down Expand Up @@ -312,7 +313,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
Expand Down Expand Up @@ -567,9 +568,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
Expand Down Expand Up @@ -901,7 +900,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
Expand Down Expand Up @@ -1073,9 +1072,7 @@ async def create(
)


async def empty(
shape: ChunkCoords, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
async def empty(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray:
"""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.

Expand All @@ -1096,9 +1093,7 @@ async def empty(
return await create(shape=shape, fill_value=None, **kwargs)


async def empty_like(
a: ArrayLike, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
async def empty_like(a: ArrayLike, **kwargs: Any) -> AnyAsyncArray:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if you have the energy for it, I think we could type **kwargs on one place with a typeddict and use Unpack in the signatures here.

this would allow us to use an overload to declare that empty_like(..., zarr_format=2) -> AsyncArray[ArrayV2Metadata]

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would like to leave anything beyond the new parametrization and using the new Any* types to follow up PRs if that's okay, to keep this PR tightly scoped.

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

Expand All @@ -1125,9 +1120,7 @@ async def empty_like(


# TODO: add type annotations for fill_value and kwargs
async def full(
shape: ChunkCoords, fill_value: Any, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
async def full(shape: ChunkCoords, fill_value: Any, **kwargs: Any) -> AnyAsyncArray:
"""Create an array, with `fill_value` being used as the default value for
uninitialized portions of the array.

Expand All @@ -1149,9 +1142,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
Expand All @@ -1172,9 +1163,7 @@ async def full_like(
return await full(**like_kwargs)


async def ones(
shape: ChunkCoords, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
async def ones(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray:
"""Create an array, with one being used as the default value for
uninitialized portions of the array.

Expand All @@ -1193,9 +1182,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
Expand All @@ -1222,7 +1209,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
Expand Down Expand Up @@ -1270,9 +1257,7 @@ async def open_array(
raise


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
Expand All @@ -1295,9 +1280,7 @@ async def open_like(
return await open_array(path=path, **like_kwargs)


async def zeros(
shape: ChunkCoords, **kwargs: Any
) -> AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata]:
async def zeros(shape: ChunkCoords, **kwargs: Any) -> AnyAsyncArray:
"""Create an array, with zero being used as the default value for
uninitialized portions of the array.

Expand All @@ -1316,9 +1299,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
Expand Down
33 changes: 17 additions & 16 deletions src/zarr/api/synchronous.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
)
from zarr.core.dtype import ZDTypeLike
from zarr.storage import StoreLike
from zarr.types import AnyArray

__all__ = [
"array",
Expand Down Expand Up @@ -169,7 +170,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
Expand Down Expand Up @@ -366,7 +367,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
Expand Down Expand Up @@ -634,7 +635,7 @@ def create(
storage_options: dict[str, Any] | None = None,
config: ArrayConfigLike | None = None,
**kwargs: Any,
) -> Array:
) -> AnyArray:
"""Create an array.

Parameters
Expand Down Expand Up @@ -770,7 +771,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`.
Expand Down Expand Up @@ -918,7 +919,7 @@ def create_array(
def from_array(
store: str | StoreLike,
*,
data: Array | npt.ArrayLike,
data: AnyArray | npt.ArrayLike,
write_data: bool = True,
name: str | None = None,
chunks: Literal["auto", "keep"] | ChunkCoords = "keep",
Expand All @@ -935,7 +936,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
Expand Down Expand Up @@ -1129,7 +1130,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.

Expand All @@ -1156,7 +1157,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.

Expand All @@ -1182,7 +1183,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
Expand All @@ -1204,7 +1205,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
Expand All @@ -1223,7 +1224,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
Expand All @@ -1242,7 +1243,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
Expand All @@ -1268,7 +1269,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
Expand Down Expand Up @@ -1304,7 +1305,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
Expand All @@ -1325,7 +1326,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
Expand All @@ -1344,7 +1345,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
Expand Down
Loading
Loading