diff --git a/changes/4366.feature.md b/changes/4366.feature.md new file mode 100644 index 0000000000..69f0d369a3 --- /dev/null +++ b/changes/4366.feature.md @@ -0,0 +1,45 @@ +`zarr.open` now opens a Zarr format 3 array or group with a single read of +`zarr.json`, where it used to make up to seven requests: it looked for an array +first and then opened a group, and the two steps read overlapping keys. A format +2 node is read in one concurrent round, and an opener that knows which kind it +wants reads only that kind's documents, so `zarr.open_array` on a format 2 array +reads `.zarray` and `.zattrs` and nothing else. + +`zarr.open` also has a stated contract now: given `shape` it behaves as +`zarr.open_array`; otherwise the reading modes (`r`, `r+`, `a`) open whatever +node is at the path, with `a` creating a group when there is none, and the +creating modes (`w`, `w-`) create a group, `w` replacing whatever is there. + +Some errors changed type as a result, so code that catches them may need +updating: + +- `zarr.open` in mode `r` or `r+` with nothing at the path now raises + `NodeNotFoundError` rather than `GroupNotFoundError`. `NodeNotFoundError` is + the base class of `GroupNotFoundError`, so `except GroupNotFoundError` no + longer catches this; catch `NodeNotFoundError` (or `FileNotFoundError`, which + still works) instead. +- Opening a node as the wrong kind now raises `ContainsArrayError` or + `ContainsGroupError` where a "not found" error was raised before: notably + `zarr.open_array` and `zarr.Array.open` on a Zarr format 3 group, which used + to raise `NodeTypeValidationError`. These are `ValueError` subclasses and not + `FileNotFoundError`, so `except FileNotFoundError` around those calls needs to + catch `ContainsGroupError` as well. Zarr format 2 is unchanged here, because + an opener no longer reads the other kind's documents at all. +- A Zarr format 3 document whose `node_type` is missing or is neither `array` + nor `group` now raises `NodeTypeValidationError` everywhere, including + `Group.__getitem__` and `Group.members`. A missing `node_type` previously + opened as a group. +- `zarr.open(mode="w-")` without `shape` creates a group instead of raising + `TypeError`, and `zarr.open(shape=...)` on a Zarr format 3 group raises + `ContainsGroupError` instead of `TypeError`. +- Requesting consolidated metadata that is not there now raises the same message + for both formats, instead of `ValueError(".zmetadata")` for format 2. + +`zarr.errors.NodeNotFoundError`, the shared base of `ArrayNotFoundError` and +`GroupNotFoundError`, is now exported from `zarr.errors` and appears in the +error documentation. + +Two smaller fixes came out of the same work: `zarr.open` and `zarr.open_array` +apply `config` to an array they open, where it was previously ignored, and +`WrapperStore` forwards `supports_consolidated_metadata` to the store it wraps +instead of always reporting `True`. diff --git a/changes/4366.removal.md b/changes/4366.removal.md new file mode 100644 index 0000000000..22720eac6e --- /dev/null +++ b/changes/4366.removal.md @@ -0,0 +1,4 @@ +The Zarr format 2 consolidated metadata key is `.zmetadata` and is no longer +configurable. `use_consolidated` is a `bool` or `None`; passing a string, which +the `zarr.open_group` and `AsyncGroup.open` docstrings described as a way to +read consolidated metadata from a non-default key, now raises `TypeError`. diff --git a/src/zarr/api/asynchronous.py b/src/zarr/api/asynchronous.py index 1fc10cdd1e..7558077696 100644 --- a/src/zarr/api/asynchronous.py +++ b/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, TypedDict import numpy as np import numpy.typing as npt @@ -17,7 +17,6 @@ CompressorLike, create_array, from_array, - get_array_metadata, ) from zarr.core.array_spec import ArrayConfigLike, parse_array_config from zarr.core.buffer import NDArrayLike @@ -35,13 +34,16 @@ AsyncGroup, ConsolidatedMetadata, GroupMetadata, + _open_array, + _open_group, + _open_node, create_hierarchy, ) -from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata +from zarr.core.metadata import ArrayV2Metadata from zarr.errors import ( ArrayNotFoundError, GroupNotFoundError, - NodeTypeValidationError, + NodeNotFoundError, ZarrDeprecationWarning, ZarrRuntimeWarning, ZarrUserWarning, @@ -356,7 +358,8 @@ async def open( (fail if exists). If the store is read-only, the default is 'r'; otherwise, it is 'a'. zarr_format : {2, 3, None}, optional - The zarr format to use when saving. + The Zarr format of the node. None opens whichever format is found, + trying Zarr format 3 first, and creates the default format. path : str or None, optional The path within the store to open. storage_options : dict @@ -377,6 +380,15 @@ async def open( Notes ----- + What `open` opens or creates follows two rules. If `shape` is given, the + call describes an array and behaves as + [`open_array`][zarr.api.asynchronous.open_array] with the same arguments. + Otherwise, in the modes that read ('r', 'r+' and 'a'), the node at `path` is + opened whichever kind it is; when there is none, 'r' and 'r+' raise + [`NodeNotFoundError`][zarr.errors.NodeNotFoundError] and 'a' creates a + group. The modes that only create ('w' and 'w-') create a group, 'w' + replacing whatever is at `path` and 'w-' failing if anything is. + `open` returns a lazy [`Array`][zarr.Array] or [`Group`][zarr.Group] backed by the store, so data is read and written incrementally. Use [`load`][zarr.load] instead when you want the data eagerly read into an in-memory array (a @@ -390,30 +402,28 @@ async def open( mode = "a" store_path = await make_store_path(store, mode=mode, path=path, storage_options=storage_options) - # TODO: the mode check below seems wrong! - if "shape" not in kwargs and mode in {"a", "r", "r+", "w"}: - try: - metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) - # TODO: remove this cast when we fix typing for array metadata dicts - _metadata_dict = cast("ArrayMetadataDict", metadata_dict) - # for v2, the above would already have raised an exception if not an array - zarr_format = _metadata_dict["zarr_format"] - is_v3_array = zarr_format == 3 and _metadata_dict.get("node_type") == "array" - if is_v3_array or zarr_format == 2: - return AsyncArray( - store_path=store_path, metadata=_metadata_dict, config=kwargs.get("config") - ) - except (FileNotFoundError, NodeTypeValidationError): - pass - return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) - - try: - return await open_array(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) - except (KeyError, NodeTypeValidationError): - # KeyError for a missing key - # NodeTypeValidationError for failing to parse node metadata as an array when it's - # actually a group - return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) + # `config` is an array's; it applies to an array opened or created here and is + # not something a group takes + config = kwargs.pop("config", None) + if "shape" in kwargs: + # the call describes an array + return await open_array( + store=store_path, zarr_format=zarr_format, mode=mode, config=config, **kwargs + ) + if mode in _READ_MODES: + node = await _open_node( + store_path, + zarr_format=zarr_format, + use_consolidated=kwargs.get("use_consolidated"), + config=config, + ) + if node is not None: + return node + if mode != "a": + msg = f"No array or group found in store {store_path.store} at path {store_path.path!r}" + raise NodeNotFoundError(msg) + # nothing to open, or a mode that only creates: make a group + return await open_group(store=store_path, zarr_format=zarr_format, mode=mode, **kwargs) async def open_consolidated( @@ -789,7 +799,7 @@ async def open_group( zarr_format: ZarrFormat | None = None, meta_array: Any | None = None, # not used attributes: dict[str, JSON] | None = None, - use_consolidated: bool | str | None = None, + use_consolidated: bool | None = None, ) -> AsyncGroup: """Open a group using file-mode-like semantics. @@ -824,7 +834,7 @@ async def open_group( to users. Use `numpy.empty(())` by default. attributes : dict A dictionary of JSON-serializable values with user-defined attributes. - use_consolidated : bool or str, default None + use_consolidated : bool, default None Whether to use consolidated metadata. By default, consolidated metadata is used if it's present in the @@ -837,10 +847,6 @@ async def open_group( To explicitly *not* use consolidated metadata, set `use_consolidated=False`, which will fall back to using the regular, non consolidated metadata. - Zarr format 2 allowed configuring the key storing the consolidated metadata - (`.zmetadata` by default). Specify the custom key as `use_consolidated` - to load consolidated metadata from a non-default key. - Returns ------- g : group @@ -857,24 +863,18 @@ async def open_group( ) store_path = await make_store_path(store, mode=mode, storage_options=storage_options, path=path) - if attributes is None: - attributes = {} - - try: - if mode in _READ_MODES: - return await AsyncGroup.open( - store_path, zarr_format=zarr_format, use_consolidated=use_consolidated - ) - except (KeyError, FileNotFoundError): - pass + if mode in _READ_MODES: + group = await _open_group( + store_path, zarr_format=zarr_format, use_consolidated=use_consolidated + ) + if group is not None: + return group if mode in _CREATE_MODES: - overwrite = _infer_overwrite(mode) - _zarr_format = zarr_format or _default_zarr_format() return await AsyncGroup.from_store( store_path, - zarr_format=_zarr_format, - overwrite=overwrite, - attributes=attributes, + zarr_format=zarr_format or _default_zarr_format(), + overwrite=_infer_overwrite(mode), + attributes=attributes or {}, ) msg = f"No group found in store {store!r} at path {store_path.path!r}" raise GroupNotFoundError(msg) @@ -1266,20 +1266,20 @@ async def open_array( if "write_empty_chunks" in kwargs: _warn_write_empty_chunks_kwarg() - try: - return await AsyncArray.open(store_path, zarr_format=zarr_format) - except FileNotFoundError as err: - if not store_path.read_only and mode in _CREATE_MODES: - overwrite = _infer_overwrite(mode) - _zarr_format = zarr_format or _default_zarr_format() - return await create( - store=store_path, - zarr_format=_zarr_format, - overwrite=overwrite, - **kwargs, - ) - msg = f"No array found in store {store_path.store} at path {store_path.path}" - raise ArrayNotFoundError(msg) from err + if mode not in _OVERWRITE_MODES: + # whatever array is here is what the caller gets + array = await _open_array(store_path, zarr_format=zarr_format, config=kwargs.get("config")) + if array is not None: + return array + if not store_path.read_only and mode in _CREATE_MODES: + return await create( + store=store_path, + zarr_format=zarr_format or _default_zarr_format(), + overwrite=_infer_overwrite(mode), + **kwargs, + ) + msg = f"No array found in store {store_path.store} at path {store_path.path}" + raise ArrayNotFoundError(msg) async def open_like(a: ArrayLike, path: str, **kwargs: Any) -> AnyAsyncArray: diff --git a/src/zarr/api/synchronous.py b/src/zarr/api/synchronous.py index 6975f6d953..75211278cf 100644 --- a/src/zarr/api/synchronous.py +++ b/src/zarr/api/synchronous.py @@ -203,7 +203,8 @@ def open( (fail if exists). If the store is read-only, the default is 'r'; otherwise, it is 'a'. zarr_format : {2, 3, None}, optional - The zarr format to use when saving. + The Zarr format of the node. None opens whichever format is found, + trying Zarr format 3 first, and creates the default format. path : str or None, optional The path within the store to open. storage_options : dict @@ -224,6 +225,14 @@ def open( Notes ----- + What `open` opens or creates follows two rules. If `shape` is given, the + call describes an array and behaves as [`open_array`][zarr.open_array] with + the same arguments. Otherwise, in the modes that read ('r', 'r+' and 'a'), + the node at `path` is opened whichever kind it is; when there is none, 'r' + and 'r+' raise [`NodeNotFoundError`][zarr.errors.NodeNotFoundError] and 'a' + creates a group. The modes that only create ('w' and 'w-') create a group, + 'w' replacing whatever is at `path` and 'w-' failing if anything is. + `open` returns a lazy [`Array`][zarr.Array] or [`Group`][zarr.Group] backed by the store, so data is read and written incrementally. Use [`load`][zarr.load] instead when you want the data eagerly read into an in-memory array (a @@ -491,7 +500,7 @@ def open_group( zarr_format: ZarrFormat | None = None, meta_array: Any | None = None, # not used in async api attributes: dict[str, JSON] | None = None, - use_consolidated: bool | str | None = None, + use_consolidated: bool | None = None, ) -> Group: """Open a group using file-mode-like semantics. @@ -526,7 +535,7 @@ def open_group( to users. Use `numpy.empty(())` by default. attributes : dict A dictionary of JSON-serializable values with user-defined attributes. - use_consolidated : bool or str, default None + use_consolidated : bool, default None Whether to use consolidated metadata. By default, consolidated metadata is used if it's present in the @@ -539,10 +548,6 @@ def open_group( To explicitly *not* use consolidated metadata, set `use_consolidated=False`, which will fall back to using the regular, non consolidated metadata. - Zarr format 2 allowed configuring the key storing the consolidated metadata - (`.zmetadata` by default). Specify the custom key as `use_consolidated` - to load consolidated metadata from a non-default key. - Returns ------- g : Group diff --git a/src/zarr/core/array.py b/src/zarr/core/array.py index 5a8d6bf57e..dc52bb01b5 100644 --- a/src/zarr/core/array.py +++ b/src/zarr/core/array.py @@ -3,7 +3,6 @@ import copy import math import warnings -from asyncio import gather from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass, field, replace from itertools import starmap @@ -29,7 +28,6 @@ from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec from zarr.codecs.zstd import ZstdCodec from zarr.core._info import ArrayInfo -from zarr.core._json import buffer_to_json_object from zarr.core.array_spec import ArrayConfig, ArrayConfigLike, ArraySpec, parse_array_config from zarr.core.attributes import Attributes from zarr.core.buffer import ( @@ -39,7 +37,6 @@ NDBuffer, default_buffer_prototype, ) -from zarr.core.buffer.cpu import buffer_prototype as cpu_buffer_prototype from zarr.core.chunk_grids import ( SHARDED_INNER_CHUNK_MAX_BYTES, ChunkGrid, @@ -59,9 +56,6 @@ ) from zarr.core.common import ( JSON, - ZARR_JSON, - ZARRAY_JSON, - ZATTRS_JSON, ChunksLike, DimensionNamesLike, MemoryOrder, @@ -130,13 +124,11 @@ RectilinearChunkGridMetadata, RegularChunkGridMetadata, create_chunk_grid_metadata, - parse_node_type_array, ) from zarr.core.sync import sync from zarr.errors import ( ArrayNotFoundError, ChunkNotFoundError, - MetadataValidationError, ZarrDeprecationWarning, ZarrUserWarning, ) @@ -158,7 +150,6 @@ from zarr.abc.codec import CodecPipeline from zarr.abc.store import Store from zarr.codecs.sharding import IndexLocation, ShardingCodec - from zarr.core.buffer import Buffer from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar from zarr.storage import StoreLike from zarr.types import AnyArray, AnyAsyncArray, ArrayV2, ArrayV3, AsyncArrayV2, AsyncArrayV3 @@ -269,69 +260,34 @@ def create_codec_pipeline(metadata: ArrayMetadata, *, store: Store | None = None raise TypeError # pragma: no cover -async def get_array_metadata( - store_path: StorePath, zarr_format: ZarrFormat | None = 3 -) -> dict[str, JSON]: - if zarr_format == 2: - zarray_bytes, zattrs_bytes = await gather( - (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype), - (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype), - ) - if zarray_bytes is None: - msg = ( - "A Zarr V2 array metadata document was not found in store " - f"{store_path.store!r} at path {store_path.path!r}." - ) - raise ArrayNotFoundError(msg) - return _array_metadata_dict_v2(zarray_bytes, zattrs_bytes) - elif zarr_format == 3: - zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype) - if zarr_json_bytes is None: - msg = ( - "A Zarr V3 array metadata document was not found in store " - f"{store_path.store!r} at path {store_path.path!r}." - ) - raise ArrayNotFoundError(msg) - return _array_metadata_dict_v3(zarr_json_bytes) - elif zarr_format is None: - zarr_json_bytes, zarray_bytes, zattrs_bytes = await gather( - (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype), - (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype), - (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype), - ) - if zarr_json_bytes is not None and zarray_bytes is not None: - # warn and favor v3 - msg = f"Both zarr.json (Zarr format 3) and .zarray (Zarr format 2) metadata objects exist at {store_path}. Zarr v3 will be used." - warnings.warn(msg, category=ZarrUserWarning, stacklevel=1) - # favor v3 when both are present - if zarr_json_bytes is not None: - return _array_metadata_dict_v3(zarr_json_bytes) - if zarray_bytes is not None: - return _array_metadata_dict_v2(zarray_bytes, zattrs_bytes) - msg = ( - f"Neither Zarr V3 nor Zarr V2 array metadata documents " - f"were found in store {store_path.store!r} at path {store_path.path!r}." - ) - raise ArrayNotFoundError(msg) - else: - msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable] - raise MetadataValidationError(msg) +async def _read_array_metadata( + store_path: StorePath, zarr_format: ZarrFormat | None +) -> ArrayMetadata: + """The array metadata at `store_path`. + Raises `ArrayNotFoundError` if there is no array there and + `ContainsGroupError` if a Zarr format 3 group is. + """ + # group.py imports this module, so the shared reader is imported here + from zarr.core.group import read_array_metadata -def _array_metadata_dict_v2(zarray_bytes: Buffer, zattrs_bytes: Buffer | None) -> dict[str, JSON]: - """Combine a `.zarray` document and an optional `.zattrs` document into one metadata dict.""" - metadata_dict: dict[str, JSON] = buffer_to_json_object(zarray_bytes) - metadata_dict["attributes"] = ( - buffer_to_json_object(zattrs_bytes) if zattrs_bytes is not None else {} - ) - return metadata_dict + metadata = await read_array_metadata(store_path.store, store_path.path, zarr_format) + if metadata is None: + msg = f"No array found in store {store_path.store} at path {store_path.path!r}" + raise ArrayNotFoundError(msg) + return metadata -def _array_metadata_dict_v3(zarr_json_bytes: Buffer) -> dict[str, JSON]: - """Parse a `zarr.json` document, checking that it describes an array.""" - metadata_dict: dict[str, JSON] = buffer_to_json_object(zarr_json_bytes) - parse_node_type_array(metadata_dict.get("node_type")) - return metadata_dict +async def get_array_metadata( + store_path: StorePath, zarr_format: ZarrFormat | None = 3 +) -> dict[str, JSON]: + """The array metadata at `store_path`, as a dict. + + The dict is the parsed metadata serialized again, so it is normalized (for + example, defaults filled in) rather than the stored document verbatim. Prefer + `AsyncArray.open`, which builds the array from the parsed metadata directly. + """ + return (await _read_array_metadata(store_path, zarr_format)).to_dict() async def _prepare_overwrite( @@ -813,10 +769,8 @@ async def example(): ``` """ store_path = await make_store_path(store) - metadata_dict = await get_array_metadata(store_path, zarr_format=zarr_format) - # TODO: remove this cast when we have better type hints - _metadata_dict = cast("ArrayMetadataJSON_V3", metadata_dict) - return cls(store_path=store_path, metadata=_metadata_dict) + metadata = await _read_array_metadata(store_path, zarr_format) + return cls(store_path=store_path, metadata=metadata) @property def store(self) -> Store: diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index d734e6b7cd..e8c8be15ab 100644 --- a/src/zarr/core/group.py +++ b/src/zarr/core/group.py @@ -85,6 +85,107 @@ logger = logging.getLogger("zarr.group") +def _resolve_use_consolidated(store: Store, use_consolidated: bool | None) -> bool | None: + """Settle `use_consolidated` against what `store` supports. + + A store that can't hold consolidated metadata makes the answer False, unless + consolidated metadata was explicitly asked for, which is an error. + """ + if not isinstance(use_consolidated, bool | None): + raise TypeError(f"use_consolidated must be a bool or None. Got {use_consolidated!r}.") + if store.supports_consolidated_metadata: + return use_consolidated + if use_consolidated: + raise ValueError( + f"The Zarr store in use ({type(store).__name__}) doesn't support consolidated metadata." + ) + return False + + +def _apply_use_consolidated( + metadata: GroupMetadata, use_consolidated: bool | None, *, store_path: StorePath +) -> GroupMetadata: + """Enforce the caller's `use_consolidated` on freshly read group metadata. + + True requires consolidated metadata to be present; False drops whatever is + present; None keeps whatever is present. + """ + if use_consolidated and metadata.consolidated_metadata is None: + msg = ( + "Consolidated metadata requested with 'use_consolidated=True' " + f"but not found in '{store_path.path}'." + ) + raise ValueError(msg) + if use_consolidated is False and metadata.consolidated_metadata is not None: + return replace(metadata, consolidated_metadata=None) + return metadata + + +async def _open_node( + store_path: StorePath, + *, + zarr_format: ZarrFormat | None, + use_consolidated: bool | None, + config: ArrayConfigLike | None = None, +) -> AnyAsyncArray | AsyncGroup | None: + """The node at `store_path`, whichever kind it is, or None if there is none. + + For a caller that does not know what kind of node to expect. It reads the + node's metadata, each document at most once, applies `use_consolidated` if + the node is a group, and builds the node. It applies no mode policy and raises nothing for a missing + node; the callers decide what those mean. + """ + store = store_path.store + # A store that can't hold consolidated metadata has none to read. Whether the + # caller may ask for it is a question for a group, so it waits until we know. + consolidated = store.supports_consolidated_metadata and use_consolidated is not False + metadata = await read_node_metadata( + store, store_path.path, zarr_format, consolidated=consolidated + ) + if metadata is None: + return None + if isinstance(metadata, GroupMetadata): + use_consolidated = _resolve_use_consolidated(store, use_consolidated) + metadata = _apply_use_consolidated(metadata, use_consolidated, store_path=store_path) + return _build_node(store=store, path=store_path.path, metadata=metadata, config=config) + + +async def _open_array( + store_path: StorePath, *, zarr_format: ZarrFormat | None, config: ArrayConfigLike | None = None +) -> AnyAsyncArray | None: + """The array at `store_path`, or None if there is none. + + Reads only what an array needs. Raises `ContainsGroupError` for a Zarr format + 3 group, whose one document says what it is; a format 2 group is not looked + for and reads as None. + """ + metadata = await read_array_metadata(store_path.store, store_path.path, zarr_format) + if metadata is None: + return None + return _build_node( + store=store_path.store, path=store_path.path, metadata=metadata, config=config + ) + + +async def _open_group( + store_path: StorePath, *, zarr_format: ZarrFormat | None, use_consolidated: bool | None +) -> AsyncGroup | None: + """The group at `store_path`, with `use_consolidated` applied, or None if there is none. + + Reads only what a group needs. Raises `ContainsArrayError` for a Zarr format + 3 array, whose one document says what it is; a format 2 array is not looked + for and reads as None. + """ + use_consolidated = _resolve_use_consolidated(store_path.store, use_consolidated) + metadata = await read_group_metadata( + store_path.store, store_path.path, zarr_format, consolidated=use_consolidated is not False + ) + if metadata is None: + return None + metadata = _apply_use_consolidated(metadata, use_consolidated, store_path=store_path) + return AsyncGroup(metadata=metadata, store_path=store_path) + + def parse_zarr_format(data: Any) -> ZarrFormat: """Parse the zarr_format field from metadata.""" return cast("ZarrFormat", parse_field(data, Literal[2, 3], "zarr_format")) @@ -498,7 +599,7 @@ async def open( cls, store: StoreLike, zarr_format: ZarrFormat | None = 3, - use_consolidated: bool | str | None = None, + use_consolidated: bool | None = None, ) -> AsyncGroup: """Open a new AsyncGroup @@ -506,7 +607,7 @@ async def open( ---------- store : StoreLike zarr_format : {2, 3}, optional - use_consolidated : bool or str, default None + use_consolidated : bool, default None Whether to use consolidated metadata. By default, consolidated metadata is used if it's present in the @@ -519,167 +620,15 @@ async def open( To explicitly *not* use consolidated metadata, set ``use_consolidated=False``, which will fall back to using the regular, non consolidated metadata. - - Zarr format 2 allowed configuring the key storing the consolidated metadata - (``.zmetadata`` by default). Specify the custom key as ``use_consolidated`` - to load consolidated metadata from a non-default key. """ store_path = await make_store_path(store) - if not store_path.store.supports_consolidated_metadata: - # Fail if consolidated metadata was requested but the Store doesn't support it - if use_consolidated: - store_name = type(store_path.store).__name__ - raise ValueError( - f"The Zarr store in use ({store_name}) doesn't support consolidated metadata." - ) - - # if use_consolidated was None (optional), the Store dictates it doesn't want consolidation - use_consolidated = False - - consolidated_key = ZMETADATA_V2_JSON - - if (zarr_format == 2 or zarr_format is None) and isinstance(use_consolidated, str): - consolidated_key = use_consolidated - - if zarr_format == 2: - paths = [store_path / ZGROUP_JSON, store_path / ZATTRS_JSON] - if use_consolidated or use_consolidated is None: - paths.append(store_path / consolidated_key) - - zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather( - *[path.get() for path in paths] - ) - if zgroup_bytes is None: - raise FileNotFoundError(store_path) - - if use_consolidated or use_consolidated is None: - maybe_consolidated_metadata_bytes = rest[0] - - else: - maybe_consolidated_metadata_bytes = None - - elif zarr_format == 3: - zarr_json_bytes = await (store_path / ZARR_JSON).get() - if zarr_json_bytes is None: - raise FileNotFoundError(store_path) - elif zarr_format is None: - ( - zarr_json_bytes, - zgroup_bytes, - zattrs_bytes, - maybe_consolidated_metadata_bytes, - ) = await asyncio.gather( - (store_path / ZARR_JSON).get(), - (store_path / ZGROUP_JSON).get(), - (store_path / ZATTRS_JSON).get(), - (store_path / str(consolidated_key)).get(), - ) - if zarr_json_bytes is not None and zgroup_bytes is not None: - # warn and favor v3 - msg = f"Both zarr.json (Zarr format 3) and .zgroup (Zarr format 2) metadata objects exist at {store_path}. Zarr format 3 will be used." - warnings.warn(msg, category=ZarrUserWarning, stacklevel=1) - if zarr_json_bytes is None and zgroup_bytes is None: - raise FileNotFoundError( - f"could not find zarr.json or .zgroup objects in {store_path}" - ) - # set zarr_format based on which keys were found - if zarr_json_bytes is not None: - zarr_format = 3 - else: - zarr_format = 2 - else: - msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable] - raise MetadataValidationError(msg) - - if zarr_format == 2: - if zgroup_bytes is None: - raise FileNotFoundError(store_path) - - if use_consolidated and maybe_consolidated_metadata_bytes is None: - # the user requested consolidated metadata, but it was missing - raise ValueError(consolidated_key) - - elif use_consolidated is False: - # the user explicitly opted out of consolidated_metadata. - # Discard anything we might have read. - maybe_consolidated_metadata_bytes = None - - return cls._from_bytes_v2( - store_path, zgroup_bytes, zattrs_bytes, maybe_consolidated_metadata_bytes - ) - else: - # V3 groups are comprised of a zarr.json object - if zarr_json_bytes is None: - raise FileNotFoundError(store_path) - if not isinstance(use_consolidated, bool | None): - raise TypeError("use_consolidated must be a bool or None for Zarr format 3.") - - return cls._from_bytes_v3( - store_path, - zarr_json_bytes, - use_consolidated=use_consolidated, - ) - - @classmethod - def _from_bytes_v2( - cls, - store_path: StorePath, - zgroup_bytes: Buffer, - zattrs_bytes: Buffer | None, - consolidated_metadata_bytes: Buffer | None, - ) -> AsyncGroup: - # V2 groups are comprised of a .zgroup and .zattrs objects - zgroup = buffer_to_json_object(zgroup_bytes) - zattrs = buffer_to_json_object(zattrs_bytes) if zattrs_bytes is not None else {} - group_metadata: dict[str, Any] = {**zgroup, "attributes": zattrs} - - if consolidated_metadata_bytes is not None: - v2_consolidated_doc = buffer_to_json_object(consolidated_metadata_bytes) - v2_consolidated_metadata = cast("dict[str, Any]", v2_consolidated_doc["metadata"]) - # We already read zattrs and zgroup. Should we ignore these? - v2_consolidated_metadata.pop(".zattrs", None) - v2_consolidated_metadata.pop(".zgroup", None) - - consolidated_metadata: defaultdict[str, dict[str, Any]] = defaultdict(dict) - - # keys like air/.zarray, air/.zattrs - for k, v in v2_consolidated_metadata.items(): - path, kind = k.rsplit("/.", 1) - - if kind == "zarray": - consolidated_metadata[path].update(v) - elif kind == "zattrs": - consolidated_metadata[path]["attributes"] = v - elif kind == "zgroup": - consolidated_metadata[path].update(v) - else: - raise ValueError(f"Invalid file type '{kind}' at path '{path}") - - group_metadata["consolidated_metadata"] = { - "metadata": dict(consolidated_metadata), - "kind": "inline", - "must_understand": False, - } - - return cls.from_dict(store_path, group_metadata) - - @classmethod - def _from_bytes_v3( - cls, - store_path: StorePath, - zarr_json_bytes: Buffer, - use_consolidated: bool | None, - ) -> AsyncGroup: - group_metadata = buffer_to_json_object(zarr_json_bytes) - if use_consolidated and group_metadata.get("consolidated_metadata") is None: - msg = f"Consolidated metadata requested with 'use_consolidated=True' but not found in '{store_path.path}'." - raise ValueError(msg) - - elif use_consolidated is False: - # Drop consolidated metadata if it's there. - group_metadata.pop("consolidated_metadata", None) - - return cls.from_dict(store_path, group_metadata) + group = await _open_group( + store_path, zarr_format=zarr_format, use_consolidated=use_consolidated + ) + if group is None: + msg = f"No group found in store {store_path.store} at path {store_path.path!r}" + raise GroupNotFoundError(msg) + return group @classmethod def from_dict( @@ -687,13 +636,10 @@ def from_dict( store_path: StorePath, data: dict[str, Any], ) -> AsyncGroup: - node_type = data.pop("node_type", None) - if node_type == "array": + if data.get("node_type") == "array": msg = f"An array already exists in store {store_path.store} at path {store_path.path}." raise ContainsArrayError(msg) - elif node_type not in ("group", None): - msg = f"Node type in metadata ({node_type}) is not 'group'" - raise GroupNotFoundError(msg) + # any other wrong node_type is GroupMetadata.from_dict's to reject return cls( metadata=GroupMetadata.from_dict(data), store_path=store_path, @@ -3493,61 +3439,261 @@ async def _iter_members_deep( yield key, node -async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata: - """ - Given a store_path, return ArrayV3Metadata or GroupMetadata defined by the metadata - document stored at store_path.path / zarr.json. If no such document is found, raise a - FileNotFoundError. +@overload +async def read_node_metadata( + store: Store, path: str, zarr_format: Literal[3], *, consolidated: bool = False +) -> ArrayV3Metadata | GroupMetadata | None: ... + + +@overload +async def read_node_metadata( + store: Store, path: str, zarr_format: Literal[2], *, consolidated: bool = False +) -> ArrayV2Metadata | GroupMetadata | None: ... + + +@overload +async def read_node_metadata( + store: Store, path: str, zarr_format: ZarrFormat | None, *, consolidated: bool = False +) -> ArrayV2Metadata | ArrayV3Metadata | GroupMetadata | None: ... + + +async def read_node_metadata( + store: Store, path: str, zarr_format: ZarrFormat | None, *, consolidated: bool = False +) -> ArrayV2Metadata | ArrayV3Metadata | GroupMetadata | None: + """Read the metadata of the node at `path`, whichever kind and format it is. + + With `zarr_format` None, Zarr format 3 is tried first and format 2 only if + there is no `zarr.json`, so a path holding both formats is read as format 3 + without a second look. `consolidated` says whether the format 2 reader also + reads the consolidated-metadata document. Returns None if no node is found. """ + if zarr_format == 3: + return await read_v3_metadata(store, path) + if zarr_format == 2: + return await read_v2_metadata(store, path, consolidated=consolidated) + if zarr_format is None: + metadata = await read_v3_metadata(store, path) + if metadata is None: + return await read_v2_metadata(store, path, consolidated=consolidated) + return metadata + msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable] + raise MetadataValidationError(msg) + + +async def _read_zarr_json(store: Store, path: str) -> dict[str, JSON] | None: + """The parsed `zarr.json` document at `path`, or None if there is none.""" zarr_json_bytes = await store.get( _join_paths([path, ZARR_JSON]), prototype=default_buffer_prototype() ) if zarr_json_bytes is None: - raise FileNotFoundError(path) - return _build_metadata_v3(buffer_to_json_object(zarr_json_bytes)) + return None + return buffer_to_json_object(zarr_json_bytes) -async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupMetadata: +async def read_v3_metadata(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata | None: + """Read the Zarr format 3 node metadata at `path`, or None if there is no `zarr.json`. + + One read: the document's `node_type` says whether it is an array or a group. + """ + doc = await _read_zarr_json(store, path) + if doc is None: + return None + return _build_metadata_v3(doc) + + +async def read_v2_metadata( + store: Store, path: str, *, consolidated: bool = False +) -> ArrayV2Metadata | GroupMetadata | None: + """Read the Zarr format 2 node metadata at `path`, not knowing which kind it is. + + Both kinds' documents are read in one concurrent round: `.zarray`, `.zgroup`, + `.zattrs` and, if `consolidated`, the `.zmetadata` document. `.zarray` makes + the node an array and `.zgroup` a group, the array winning if both exist. A + caller that knows which kind it wants should use `read_v2_array_metadata` or + `read_v2_group_metadata`, which read only that kind's documents. Returns None + if there is neither. """ - Given a store_path, return ArrayV2Metadata or GroupMetadata defined by the metadata - document stored at store_path.path / (.zgroup | .zarray). If no such document is found, - raise a FileNotFoundError. + keys = [ZARRAY_JSON, ZGROUP_JSON, ZATTRS_JSON] + if consolidated: + keys.append(ZMETADATA_V2_JSON) + zarray_bytes, zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather( + *(store.get(_join_paths([path, key]), prototype=default_buffer_prototype()) for key in keys) + ) + array = _v2_array_metadata(zarray_bytes, zattrs_bytes) + if array is not None: + return array + return _v2_group_metadata(zgroup_bytes, zattrs_bytes, rest[0] if rest else None) + + +async def read_v2_array_metadata(store: Store, path: str) -> ArrayV2Metadata | None: + """Read the Zarr format 2 array metadata at `path`, or None if there is no `.zarray`. + + One concurrent read of `.zarray` and `.zattrs`; nothing else is looked at. """ - # TODO: consider first fetching array metadata, and only fetching group metadata when we don't - # find an array - zarray_bytes, zgroup_bytes, zattrs_bytes = await asyncio.gather( + zarray_bytes, zattrs_bytes = await asyncio.gather( store.get(_join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()), - store.get(_join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()), store.get(_join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()), ) + return _v2_array_metadata(zarray_bytes, zattrs_bytes) - zattrs: dict[str, JSON] - if zattrs_bytes is None: - zattrs = {} - else: - zattrs = buffer_to_json_object(zattrs_bytes) - # TODO: decide how to handle finding both array and group metadata. The spec does not seem to - # consider this situation. A practical approach would be to ignore that combination, and only - # return the array metadata. - if zarray_bytes is not None: - zmeta = buffer_to_json_object(zarray_bytes) - else: - if zgroup_bytes is None: - # neither .zarray or .zgroup were found results in KeyError - raise FileNotFoundError(path) +async def read_v2_group_metadata( + store: Store, path: str, *, consolidated: bool = False +) -> GroupMetadata | None: + """Read the Zarr format 2 group metadata at `path`, or None if there is no `.zgroup`. + + One concurrent read of `.zgroup`, `.zattrs` and, if `consolidated`, the + `.zmetadata` document, which is attached to the group; nothing else is + looked at. + """ + keys = [ZGROUP_JSON, ZATTRS_JSON] + if consolidated: + keys.append(ZMETADATA_V2_JSON) + zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather( + *(store.get(_join_paths([path, key]), prototype=default_buffer_prototype()) for key in keys) + ) + return _v2_group_metadata(zgroup_bytes, zattrs_bytes, rest[0] if rest else None) + + +def _v2_attributes(zattrs_bytes: Buffer | None) -> dict[str, JSON]: + return {} if zattrs_bytes is None else buffer_to_json_object(zattrs_bytes) + + +def _v2_array_metadata( + zarray_bytes: Buffer | None, zattrs_bytes: Buffer | None +) -> ArrayV2Metadata | None: + """Array metadata from a `.zarray` document and its `.zattrs`, or None without the former.""" + if zarray_bytes is None: + return None + return ArrayV2Metadata.from_dict( + buffer_to_json_object(zarray_bytes) | {"attributes": _v2_attributes(zattrs_bytes)} + ) + + +def _v2_group_metadata( + zgroup_bytes: Buffer | None, zattrs_bytes: Buffer | None, consolidated_bytes: Buffer | None +) -> GroupMetadata | None: + """Group metadata from a `.zgroup` document, its `.zattrs` and an optional consolidated document.""" + if zgroup_bytes is None: + return None + metadata = GroupMetadata.from_dict( + buffer_to_json_object(zgroup_bytes) | {"attributes": _v2_attributes(zattrs_bytes)} + ) + if consolidated_bytes is None: + return metadata + consolidated = _consolidated_metadata_from_v2_doc(buffer_to_json_object(consolidated_bytes)) + return replace(metadata, consolidated_metadata=consolidated) + + +async def read_array_metadata( + store: Store, path: str, zarr_format: ZarrFormat | None +) -> ArrayV2Metadata | ArrayV3Metadata | None: + """Read the array metadata at `path`, reading only what an array needs. + + Composed like `read_node_metadata`, but the format 2 step reads only the + array documents, so a format 2 group reads as None. A format 3 group is + reported with `ContainsGroupError`, since its one document says what it is. + """ + if zarr_format not in (2, 3, None): + msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." + raise MetadataValidationError(msg) + if zarr_format != 2: + doc = await _read_zarr_json(store, path) + if doc is not None: + # say what is here before building it, so a broken group document + # still reads as "a group is here" + if doc.get("node_type") == "group": + raise ContainsGroupError(f"A group exists in store {store} at path {path!r}.") + metadata = _build_metadata_v3(doc) + if isinstance(metadata, GroupMetadata): # pragma: no cover + raise ContainsGroupError(f"A group exists in store {store} at path {path!r}.") + return metadata + if zarr_format == 3: + return None + return await read_v2_array_metadata(store, path) + + +async def read_group_metadata( + store: Store, path: str, zarr_format: ZarrFormat | None, *, consolidated: bool = False +) -> GroupMetadata | None: + """Read the group metadata at `path`, reading only what a group needs. + + Composed like `read_node_metadata`, but the format 2 step reads only the + group documents, so a format 2 array reads as None. A format 3 array is + reported with `ContainsArrayError`, since its one document says what it is. + """ + if zarr_format not in (2, 3, None): + msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." + raise MetadataValidationError(msg) + if zarr_format != 2: + doc = await _read_zarr_json(store, path) + if doc is not None: + # say what is here before building it, so a broken array document + # still reads as "an array is here" + if doc.get("node_type") == "array": + raise ContainsArrayError(f"An array exists in store {store} at path {path!r}.") + metadata = _build_metadata_v3(doc) + if isinstance(metadata, ArrayV3Metadata): # pragma: no cover + raise ContainsArrayError(f"An array exists in store {store} at path {path!r}.") + return metadata + if zarr_format == 3: + return None + return await read_v2_group_metadata(store, path, consolidated=consolidated) + + +def _consolidated_metadata_from_v2_doc(doc: dict[str, JSON]) -> ConsolidatedMetadata: + """Turn a Zarr format 2 consolidated-metadata document into `ConsolidatedMetadata`. + + The format 2 document is flat, keyed like `air/.zarray` and `air/.zattrs`; + those become one metadata dict per path. The root's own `.zgroup` and + `.zattrs` are dropped, since they were read directly. + """ + v2_metadata = doc.get("metadata") + if not isinstance(v2_metadata, dict): + msg = f"A consolidated metadata document needs a 'metadata' object. Got {doc!r}." + raise MetadataValidationError(msg) + v2_metadata = cast("dict[str, dict[str, Any]]", dict(v2_metadata)) + v2_metadata.pop(".zattrs", None) + v2_metadata.pop(".zgroup", None) + + consolidated_metadata: defaultdict[str, dict[str, Any]] = defaultdict(dict) + for k, v in v2_metadata.items(): + path, kind = k.rsplit("/.", 1) + if kind == "zarray": + consolidated_metadata[path].update(v) + elif kind == "zattrs": + consolidated_metadata[path]["attributes"] = v + elif kind == "zgroup": + consolidated_metadata[path].update(v) else: - zmeta = buffer_to_json_object(zgroup_bytes) + raise ValueError(f"Invalid file type '{kind}' at path '{path}") + return ConsolidatedMetadata.from_dict( + {"metadata": dict(consolidated_metadata), "kind": "inline", "must_understand": False} + ) - return _build_metadata_v2(zmeta, zattrs) + +async def _read_metadata_v3(store: Store, path: str) -> ArrayV3Metadata | GroupMetadata: + """`read_v3_metadata`, raising FileNotFoundError instead of returning None.""" + metadata = await read_v3_metadata(store, path) + if metadata is None: + raise FileNotFoundError(path) + return metadata + + +async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupMetadata: + """`read_v2_metadata`, raising FileNotFoundError instead of returning None.""" + metadata = await read_v2_metadata(store, path) + if metadata is None: + raise FileNotFoundError(path) + return metadata async def _read_group_metadata_v2(store: Store, path: str) -> GroupMetadata: """ Read group metadata or error """ - meta = await _read_metadata_v2(store=store, path=path) - if not isinstance(meta, GroupMetadata): + meta = await read_v2_group_metadata(store, path) + if meta is None: raise FileNotFoundError(f"Group metadata was not found in {store} at {path}") return meta @@ -3576,98 +3722,61 @@ def _build_metadata_v3(zarr_json: dict[str, JSON]) -> ArrayV3Metadata | GroupMet """ if "node_type" not in zarr_json: msg = "Required key 'node_type' is missing from the provided metadata document." - raise MetadataValidationError(msg) + raise NodeTypeValidationError(msg) match zarr_json: case {"node_type": "array"}: return ArrayV3Metadata.from_dict(zarr_json) case {"node_type": "group"}: return GroupMetadata.from_dict(zarr_json) + case {"node_type": node_type}: + msg = f"Invalid value for 'node_type'. Expected 'array' or 'group'. Got {node_type!r}." + raise NodeTypeValidationError(msg) case _: # pragma: no cover - raise ValueError( - "invalid value for `node_type` key in metadata document" - ) # pragma: no cover - - -def _build_metadata_v2( - zarr_json: dict[str, JSON], attrs_json: dict[str, JSON] -) -> ArrayV2Metadata | GroupMetadata: - """ - Convert a dict representation of Zarr V2 metadata into the corresponding metadata class. - """ - match zarr_json: - case {"shape": _}: - return ArrayV2Metadata.from_dict(zarr_json | {"attributes": attrs_json}) - case _: # pragma: no cover - return GroupMetadata.from_dict(zarr_json | {"attributes": attrs_json}) + raise AssertionError("unreachable") # pragma: no cover @overload -def _build_node(*, store: Store, path: str, metadata: ArrayV2Metadata) -> AsyncArrayV2: ... +def _build_node( + *, store: Store, path: str, metadata: ArrayV2Metadata, config: ArrayConfigLike | None = None +) -> AsyncArrayV2: ... @overload -def _build_node(*, store: Store, path: str, metadata: ArrayV3Metadata) -> AsyncArrayV3: ... +def _build_node( + *, store: Store, path: str, metadata: ArrayV3Metadata, config: ArrayConfigLike | None = None +) -> AsyncArrayV3: ... @overload -def _build_node(*, store: Store, path: str, metadata: GroupMetadata) -> AsyncGroup: ... +def _build_node( + *, store: Store, path: str, metadata: GroupMetadata, config: ArrayConfigLike | None = None +) -> AsyncGroup: ... def _build_node( - *, store: Store, path: str, metadata: ArrayV3Metadata | ArrayV2Metadata | GroupMetadata + *, + store: Store, + path: str, + metadata: ArrayV3Metadata | ArrayV2Metadata | GroupMetadata, + config: ArrayConfigLike | None = None, ) -> AnyAsyncArray | AsyncGroup: """ - Take a metadata object and return a node (AsyncArray or AsyncGroup). + Take a metadata object and return a node (AsyncArray or AsyncGroup). `config` + applies to an array and is ignored for a group. """ store_path = StorePath(store=store, path=path) match metadata: case ArrayV2Metadata() | ArrayV3Metadata(): - return AsyncArray(metadata, store_path=store_path) + return AsyncArray(metadata, store_path=store_path, config=config) case GroupMetadata(): return AsyncGroup(metadata, store_path=store_path) case _: # pragma: no cover raise ValueError(f"Unexpected metadata type: {type(metadata)}") # pragma: no cover -async def _get_node_v2(store: Store, path: str) -> AsyncArrayV2 | AsyncGroup: - """ - Read a Zarr v2 AsyncArray or AsyncGroup from a path in a Store. - - Parameters - ---------- - store : Store - The store-like object to read from. - path : str - The path to the node to read. - - Returns - ------- - AsyncArray | AsyncGroup - """ - metadata = await _read_metadata_v2(store=store, path=path) - return _build_node(store=store, path=path, metadata=metadata) - - -async def _get_node_v3(store: Store, path: str) -> AsyncArrayV3 | AsyncGroup: - """ - Read a Zarr v3 AsyncArray or AsyncGroup from a path in a Store. - - Parameters - ---------- - store : Store - The store-like object to read from. - path : str - The path to the node to read. - - Returns - ------- - AsyncArray | AsyncGroup - """ - metadata = await _read_metadata_v3(store=store, path=path) - return _build_node(store=store, path=path, metadata=metadata) - - -async def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyAsyncArray | AsyncGroup: +async def get_node( + store: Store, path: str, zarr_format: ZarrFormat | None +) -> AnyAsyncArray | AsyncGroup: """ Get an AsyncArray or AsyncGroup from a path in a Store. @@ -3677,21 +3786,17 @@ async def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyAsync The store-like object to read from. path : str The path to the node to read. - zarr_format : {2, 3} - The zarr format of the node to read. + zarr_format : {2, 3, None} + The zarr format of the node to read, or None to detect it. Returns ------- AsyncArray | AsyncGroup """ - - match zarr_format: - case 2: - return await _get_node_v2(store=store, path=path) - case 3: - return await _get_node_v3(store=store, path=path) - case _: # pragma: no cover - raise ValueError(f"Unexpected zarr format: {zarr_format}") # pragma: no cover + metadata = await read_node_metadata(store, path, zarr_format) + if metadata is None: + raise FileNotFoundError(path) + return _build_node(store=store, path=path, metadata=metadata) async def _set_return_key( diff --git a/src/zarr/core/sync_group.py b/src/zarr/core/sync_group.py index 8af514e938..c33f65c246 100644 --- a/src/zarr/core/sync_group.py +++ b/src/zarr/core/sync_group.py @@ -142,7 +142,7 @@ def create_rooted_hierarchy( return _parse_async_node(async_node) -def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyArray | Group: +def get_node(store: Store, path: str, zarr_format: ZarrFormat | None) -> AnyArray | Group: """ Get an Array or Group from a path in a Store. @@ -152,8 +152,8 @@ def get_node(store: Store, path: str, zarr_format: ZarrFormat) -> AnyArray | Gro The store-like object to read from. path : str The path to the node to read. - zarr_format : {2, 3} - The zarr format of the node to read. + zarr_format : {2, 3, None} + The zarr format of the node to read, or None to detect it. Returns ------- diff --git a/src/zarr/errors.py b/src/zarr/errors.py index 3e445de2e9..582ad8e775 100644 --- a/src/zarr/errors.py +++ b/src/zarr/errors.py @@ -11,6 +11,7 @@ "GroupNotFoundError", "MetadataValidationError", "NegativeStepError", + "NodeNotFoundError", "NodeTypeValidationError", "UnknownCodecError", "UnstableSpecificationWarning", diff --git a/src/zarr/storage/_wrapper.py b/src/zarr/storage/_wrapper.py index 6f498a655d..57f824c91a 100644 --- a/src/zarr/storage/_wrapper.py +++ b/src/zarr/storage/_wrapper.py @@ -196,6 +196,10 @@ async def delete(self, key: str) -> None: def supports_listing(self) -> bool: return self._store.supports_listing + @property + def supports_consolidated_metadata(self) -> bool: + return self._store.supports_consolidated_metadata + def list(self) -> AsyncIterator[str]: return self._store.list() diff --git a/tests/test_api.py b/tests/test_api.py index 45d0c0dee4..3daa481e56 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,7 +1,7 @@ from __future__ import annotations +import collections import inspect -import re from typing import TYPE_CHECKING, Any import zarr.codecs @@ -13,9 +13,11 @@ if TYPE_CHECKING: from collections.abc import Callable from pathlib import Path + from typing import Self - from zarr.abc.store import Store - from zarr.core.common import JSON, MemoryOrder, ZarrFormat + from zarr.abc.store import ByteRequest + from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.common import JSON, AccessModeLiteral, MemoryOrder, ZarrFormat from zarr.types import AnyArray import contextlib @@ -30,6 +32,7 @@ import zarr.api.synchronous import zarr.core.group from zarr import Array, Group +from zarr.abc.store import Store from zarr.api.synchronous import ( create, create_array, @@ -42,15 +45,21 @@ save_array, save_group, ) -from zarr.core.buffer import NDArrayLike +from zarr.core.buffer import NDArrayLike, cpu from zarr.errors import ( ArrayNotFoundError, + ContainsArrayError, + ContainsGroupError, + GroupNotFoundError, MetadataValidationError, + NodeNotFoundError, + NodeTypeValidationError, ZarrDeprecationWarning, ZarrUserWarning, ) from zarr.storage import MemoryStore from zarr.storage._utils import normalize_path +from zarr.storage._wrapper import WrapperStore from zarr.testing.utils import gpu_test @@ -360,9 +369,9 @@ def test_array_open_array_not_found_sync() -> None: def test_v2_and_v3_exist_at_same_path(store: Store) -> None: zarr.create_array(store, shape=(10,), dtype="uint8", zarr_format=3) zarr.create_array(store, shape=(10,), dtype="uint8", zarr_format=2) - msg = f"Both zarr.json (Zarr format 3) and .zarray (Zarr format 2) metadata objects exist at {store}. Zarr v3 will be used." - with pytest.warns(ZarrUserWarning, match=re.escape(msg)): - zarr.open(store=store) + # only zarr.json is read, and the format 3 node is taken without a second look + assert zarr.open(store=store).metadata.zarr_format == 3 + assert zarr.open_array(store=store).metadata.zarr_format == 3 @pytest.mark.parametrize("store", ["memory"], indirect=True) @@ -1377,6 +1386,358 @@ async def test_open_falls_back_to_open_group_async(zarr_format: ZarrFormat) -> N assert group.attrs == {"key": "value"} +class _CountingStore(WrapperStore[Store]): + """A store that records the key of every `get` it forwards.""" + + get_counts: collections.Counter[str] + + def __init__(self, store: Store) -> None: + super().__init__(store) + self.get_counts = collections.Counter() + + def _with_store(self, store: Store) -> Self: + # `_with_store` is how a store is re-made read-only, so the copy has to + # keep counting into the same tally. + new = type(self)(store) + new.get_counts = self.get_counts + return new + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + self.get_counts[key] += 1 + return await self._store.get(key, prototype, byte_range) + + +@pytest.mark.filterwarnings("ignore:Consolidated metadata") +@pytest.mark.parametrize("node", ["array", "group"]) +@pytest.mark.parametrize( + ("zarr_format", "use_consolidated"), + [(2, None), (2, False), (2, True), (3, None), (3, False), (3, True)], +) +@pytest.mark.parametrize("mode", ["r", "r+", "a"]) +@pytest.mark.parametrize("path", ["", "parent/child"]) +async def test_open_reads_only_what_the_node_needs( + node: Literal["array", "group"], + zarr_format: ZarrFormat, + use_consolidated: bool | None, + mode: AccessModeLiteral, + path: str, +) -> None: + """`open` reads a format 3 node with one request, and a format 2 node with `zarr.json` and one round. + + Whatever it finds has to be what `open_array` or `open_group` would have + returned: the same format, path, attributes and read-only-ness, and for a + group the same consolidated metadata. + """ + store = _CountingStore(MemoryStore()) + prefix = f"{path}/" if path else "" + if node == "array": + await zarr.api.asynchronous.create_array( + store, + name=path or None, + shape=(3,), + dtype="uint8", + attributes={"key": "value"}, + zarr_format=zarr_format, + ) + else: + await zarr.api.asynchronous.open_group( + store, path=path, attributes={"key": "value"}, zarr_format=zarr_format + ) + if use_consolidated: + await zarr.api.asynchronous.consolidate_metadata(store, path=path) + + store.get_counts.clear() + result = await zarr.api.asynchronous.open( + store=store, path=path, mode=mode, use_consolidated=use_consolidated + ) + assert isinstance(result, AsyncArray if node == "array" else zarr.core.group.AsyncGroup) + assert result.metadata.zarr_format == zarr_format + assert result.path == path + assert result.attrs == {"key": "value"} + assert result.store.read_only == (mode == "r") + if isinstance(result, zarr.core.group.AsyncGroup): + assert (result.metadata.consolidated_metadata is not None) == bool(use_consolidated) + + expected_keys = {"zarr.json"} + if zarr_format == 2: + expected_keys |= {".zarray", ".zgroup", ".zattrs"} + if use_consolidated is not False: + expected_keys.add(".zmetadata") + assert store.get_counts == {prefix + key: 1 for key in expected_keys} + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_open_array_does_not_fall_back(zarr_format: ZarrFormat) -> None: + """`open` on an array still returns the array rather than falling back to a group.""" + store = MemoryStore() + await zarr.api.asynchronous.create_array( + store, shape=(10,), dtype="uint8", zarr_format=zarr_format, attributes={"k": "v"} + ) + arr = await zarr.api.asynchronous.open(store=store) + assert isinstance(arr, AsyncArray) + assert arr.metadata.zarr_format == zarr_format + assert arr.attrs == {"k": "v"} + + +async def test_open_invalid_zarr_format_raises() -> None: + """An invalid `zarr_format` is a bad request, not a missing array, so it still raises.""" + store = MemoryStore() + with pytest.raises( + MetadataValidationError, + match="Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '3.0'.", + ): + await zarr.api.asynchronous.open(store=store, zarr_format="3.0") # type: ignore[arg-type] + + +async def test_open_zarr_json_without_node_type_raises() -> None: + """A format 3 document must say what it is: no `node_type` is a validation error.""" + store = MemoryStore() + await store.set( + "zarr.json", cpu.Buffer.from_bytes(b'{"zarr_format": 3, "attributes": {"k": "v"}}') + ) + with pytest.raises(NodeTypeValidationError, match="Required key 'node_type' is missing"): + await zarr.api.asynchronous.open(store=store, mode="r") + + +async def test_open_zarr_json_with_invalid_node_type_raises() -> None: + """A `node_type` that is neither array nor group is a validation error, not a missing node.""" + store = MemoryStore() + await store.set("zarr.json", cpu.Buffer.from_bytes(b'{"zarr_format": 3, "node_type": "foo"}')) + with pytest.raises(NodeTypeValidationError, match="Expected 'array' or 'group'. Got 'foo'"): + await zarr.api.asynchronous.open(store=store, mode="r") + + +# Zarr format 3 says what a node is in the one document every opener reads, so opening +# the wrong kind is reported as such. Zarr format 2 has a document per kind, and an +# opener that knows what it wants reads only that kind's, so the other kind reads as +# nothing in the read modes; the create modes still find it when they check before +# writing. +_WRONG_KIND_ERROR: dict[tuple[str, ZarrFormat], type[Exception]] = { + ("array as group", 3): ContainsArrayError, + ("array as group", 2): GroupNotFoundError, + ("group as array", 3): ContainsGroupError, + ("group as array", 2): ArrayNotFoundError, +} + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_async_array_open_on_group_raises(zarr_format: ZarrFormat) -> None: + store = MemoryStore() + await zarr.api.asynchronous.open_group(store, zarr_format=zarr_format) + with pytest.raises(_WRONG_KIND_ERROR["group as array", zarr_format]): + await AsyncArray.open(store, zarr_format=zarr_format) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_async_group_open_on_array_raises(zarr_format: ZarrFormat) -> None: + store = MemoryStore() + await zarr.api.asynchronous.create_array( + store, shape=(3,), dtype="uint8", zarr_format=zarr_format + ) + with pytest.raises(_WRONG_KIND_ERROR["array as group", zarr_format]): + await zarr.core.group.AsyncGroup.open(store, zarr_format=zarr_format) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("mode", ["r", "a"]) +def test_open_group_on_array_raises(mode: AccessModeLiteral, zarr_format: ZarrFormat) -> None: + store = MemoryStore() + zarr.create_array(store, shape=(3,), dtype="uint8", zarr_format=zarr_format) + expected = _WRONG_KIND_ERROR["array as group", zarr_format] + if mode == "a" and zarr_format == 2: + expected = ContainsArrayError # found by the create step's check + with pytest.raises(expected): + zarr.open_group(store, mode=mode, zarr_format=zarr_format) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("mode", ["r", "a"]) +def test_open_array_on_group_raises(mode: AccessModeLiteral, zarr_format: ZarrFormat) -> None: + store = MemoryStore() + zarr.create_group(store, zarr_format=zarr_format) + expected = _WRONG_KIND_ERROR["group as array", zarr_format] + if mode == "a" and zarr_format == 2: + expected = ContainsGroupError # found by the create step's check + with pytest.raises(expected): + zarr.open_array(store, mode=mode, zarr_format=zarr_format, shape=(3,), dtype="uint8") + + +@pytest.mark.parametrize("node", ["array", "group"]) +@pytest.mark.parametrize("zarr_format", [2, 3]) +@pytest.mark.parametrize("given_format", [True, False], ids=["format given", "format detected"]) +async def test_open_array_and_open_group_read_only_their_kind( + node: Literal["array", "group"], zarr_format: ZarrFormat, given_format: bool +) -> None: + """An opener that knows which kind it wants reads only that kind's documents. + + A format 3 node is one read either way. A format 2 array is `.zarray` and + `.zattrs`; a format 2 group is `.zgroup`, `.zattrs` and `.zmetadata`. Detecting + the format adds the one read of `zarr.json` that finds nothing. + """ + store = _CountingStore(MemoryStore()) + if node == "array": + await zarr.api.asynchronous.create_array( + store, shape=(3,), dtype="uint8", zarr_format=zarr_format + ) + else: + await zarr.api.asynchronous.open_group(store, zarr_format=zarr_format) + store.get_counts.clear() + + fmt = zarr_format if given_format else None + if node == "array": + await zarr.api.asynchronous.open_array(store=store, mode="r", zarr_format=fmt) + else: + await zarr.api.asynchronous.open_group(store=store, mode="r", zarr_format=fmt) + + if zarr_format == 3: + expected_keys = {"zarr.json"} + elif node == "array": + expected_keys = {".zarray", ".zattrs"} + else: + expected_keys = {".zgroup", ".zattrs", ".zmetadata"} + if zarr_format == 2 and not given_format: + expected_keys.add("zarr.json") + assert store.get_counts == dict.fromkeys(expected_keys, 1) + + +class _NoConsolidatedMemoryStore(MemoryStore): + """A store that says it cannot hold consolidated metadata.""" + + @property + def supports_consolidated_metadata(self) -> bool: + return False + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_open_use_consolidated_on_array_in_store_without_support(zarr_format: ZarrFormat) -> None: + """Asking for consolidated metadata is a question for a group; an array is returned regardless.""" + store = _NoConsolidatedMemoryStore() + zarr.create_array(store, shape=(3,), dtype="uint8", zarr_format=zarr_format) + assert isinstance(zarr.open(store=store, mode="r", use_consolidated=True), Array) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_open_use_consolidated_on_group_in_store_without_support_raises( + zarr_format: ZarrFormat, +) -> None: + store = _NoConsolidatedMemoryStore() + zarr.create_group(store, zarr_format=zarr_format) + with pytest.raises(ValueError, match="doesn't support consolidated metadata"): + zarr.open(store=store, mode="r", use_consolidated=True) + + +@pytest.mark.parametrize("zarr_format", [2, 3]) +def test_open_group_rejects_a_str_use_consolidated(zarr_format: ZarrFormat) -> None: + """The consolidated-metadata key is not configurable; `use_consolidated` is a bool or None.""" + store = MemoryStore() + zarr.create_group(store, zarr_format=zarr_format) + with pytest.raises(TypeError, match="use_consolidated must be a bool or None"): + zarr.open_group(store, mode="r", zarr_format=zarr_format, use_consolidated="custom") # type: ignore[arg-type] + + +def test_open_config_applies_to_arrays_only() -> None: + """`config` is an array's: it is applied to an array and ignored for a group, found or created.""" + store = MemoryStore() + zarr.create_array(store, name="a", shape=(3,), dtype="uint8") + zarr.create_group(store, path="g") + array = zarr.open(store=store, path="a", mode="r", config={"order": "F"}) + assert isinstance(array, Array) + assert array._async_array.config.order == "F" + assert isinstance(zarr.open(store=store, path="g", mode="r", config={"order": "F"}), Group) + assert isinstance(zarr.open(store=store, path="new", mode="a", config={"order": "F"}), Group) + + +def test_open_array_applies_config() -> None: + store = MemoryStore() + zarr.create_array(store, shape=(3,), dtype="uint8") + assert zarr.open_array(store, mode="r", config={"order": "F"})._async_array.config.order == "F" + + +async def test_wrong_kind_is_reported_before_the_document_is_built() -> None: + """A format 3 document says what it is in `node_type`; a broken document of the other kind still says so.""" + store = MemoryStore() + await store.set("zarr.json", cpu.Buffer.from_bytes(b'{"zarr_format": 3, "node_type": "array"}')) + with pytest.raises(ContainsArrayError): + await zarr.api.asynchronous.open_group(store, mode="r") + await store.set( + "zarr.json", + cpu.Buffer.from_bytes(b'{"zarr_format": 3, "node_type": "group", "shape": [1]}'), + ) + with pytest.raises(ContainsGroupError): + await zarr.api.asynchronous.open_array(store=store, mode="r") + + +def test_open_group_missing_v2_consolidated_metadata_message() -> None: + store = MemoryStore() + zarr.create_group(store, zarr_format=2) + with pytest.raises( + ValueError, match="Consolidated metadata requested with 'use_consolidated=True'" + ): + zarr.open_group(store, mode="r", zarr_format=2, use_consolidated=True) + + +async def test_open_group_malformed_v2_consolidated_document_raises() -> None: + store = MemoryStore() + await zarr.api.asynchronous.create_group(store=store, zarr_format=2) + await store.set(".zmetadata", cpu.Buffer.from_bytes(b'{"zarr_consolidated_format": 1}')) + with pytest.raises(MetadataValidationError, match="needs a 'metadata' object"): + await zarr.api.asynchronous.open_group(store, mode="r", zarr_format=2) + + +def _expected_open_outcome( + existing: str, mode: str, shape: tuple[int, ...] | None +) -> type[Array[Any] | Group | Exception]: + """The contract of `open`, as a table: what comes back for what is there, the mode, and `shape`.""" + if mode == "w-" and existing != "nothing": + return FileExistsError + if mode == "w": + return Array if shape else Group + if existing == "nothing": + if mode in ("r", "r+"): + return ArrayNotFoundError if shape else NodeNotFoundError + return Array if shape else Group + if shape: + return Array if existing == "array" else ContainsGroupError + return Array if existing == "array" else Group + + +@pytest.mark.parametrize("existing", ["nothing", "array", "group"]) +@pytest.mark.parametrize("mode", ["r", "r+", "a", "w", "w-"]) +@pytest.mark.parametrize("shape", [None, (3,)], ids=["no shape", "shape"]) +def test_open_mode_contract( + existing: str, mode: AccessModeLiteral, shape: tuple[int, ...] | None +) -> None: + """`open` follows its two rules for every mode, whatever is at the path, with and without `shape`. + + This is the contract in the default format; `tests/test_open_properties.py` + covers both formats. With `shape` the call describes an array. Without it, the reading modes open + whatever node is there and 'a' creates a group when there is none; the + creating modes make a group, 'w' over whatever is there and 'w-' only over + nothing. An opened node keeps its attributes; a created one has none. + """ + store = MemoryStore() + if existing == "array": + zarr.create_array(store, shape=(3,), dtype="uint8", attributes={"old": True}) + elif existing == "group": + zarr.create_group(store, attributes={"old": True}) + kwargs: dict[str, Any] = {} if shape is None else {"shape": shape, "dtype": "uint8"} + expected = _expected_open_outcome(existing, mode, shape) + + if issubclass(expected, Exception): + with pytest.raises(expected): + zarr.open(store=store, mode=mode, **kwargs) + return + node = zarr.open(store=store, mode=mode, **kwargs) + assert isinstance(node, expected) + opened = existing != "nothing" and mode in ("r", "r+", "a") + assert node.attrs.get("old") is (True if opened else None) + + @pytest.mark.parametrize("mode", ["r", "r+", "w", "a"]) def test_open_modes_creates_group(tmp_path: Path, mode: str) -> None: # https://github.com/zarr-developers/zarr-python/issues/2490 diff --git a/tests/test_api/test_asynchronous.py b/tests/test_api/test_asynchronous.py index 6ebec36bbd..8ef9ad0588 100644 --- a/tests/test_api/test_asynchronous.py +++ b/tests/test_api/test_asynchronous.py @@ -11,6 +11,7 @@ from zarr.api.asynchronous import _get_shape_chunks, _like_args, group, open from zarr.core.buffer.core import default_buffer_prototype from zarr.core.group import AsyncGroup +from zarr.errors import ContainsGroupError if TYPE_CHECKING: from pathlib import Path @@ -96,19 +97,17 @@ def test_like_args( assert _like_args(observed) == expected -async def test_open_no_array() -> None: +async def test_open_with_shape_on_group_raises() -> None: """ - Test that zarr.api.asynchronous.open attempts to open a group when no array is found, but shape was specified in kwargs. - This behavior makes no sense but we should still test it. + With `shape` given, `open` describes an array, so a group at the path is an error rather than + something to fall back to. """ store = { "zarr.json": default_buffer_prototype().buffer.from_bytes( json.dumps({"zarr_format": 3, "node_type": "group"}).encode("utf-8") ) } - with pytest.raises( - TypeError, match=r"open_group\(\) got an unexpected keyword argument 'shape'" - ): + with pytest.raises(ContainsGroupError, match="A group exists in store"): await open(store=store, shape=(1,)) diff --git a/tests/test_group.py b/tests/test_group.py index 31fbd138cd..5297276a7a 100644 --- a/tests/test_group.py +++ b/tests/test_group.py @@ -1547,6 +1547,7 @@ def test_open_mutable_mapping_sync(): async def test_open_ambiguous_node(): + """A path holding both formats opens as Zarr format 3, without reading or warning about the other.""" zarr_json_bytes = default_buffer_prototype().buffer.from_bytes( json.dumps({"zarr_format": 3, "node_type": "group"}).encode("utf-8") ) @@ -1554,11 +1555,8 @@ async def test_open_ambiguous_node(): json.dumps({"zarr_format": 2}).encode("utf-8") ) store: dict[str, Buffer] = {"zarr.json": zarr_json_bytes, ".zgroup": zgroup_bytes} - with pytest.warns( - ZarrUserWarning, - match=r"Both zarr\.json \(Zarr format 3\) and \.zgroup \(Zarr format 2\) metadata objects exist at", - ): - await AsyncGroup.open(store, zarr_format=None) + group = await AsyncGroup.open(store, zarr_format=None) + assert group.metadata.zarr_format == 3 class TestConsolidated: diff --git a/tests/test_metadata/test_consolidated.py b/tests/test_metadata/test_consolidated.py index cd0fd92d74..8f884a0383 100644 --- a/tests/test_metadata/test_consolidated.py +++ b/tests/test_metadata/test_consolidated.py @@ -597,16 +597,16 @@ async def test_open_consolidated_raises_async(self, zarr_format: ZarrFormat) -> async def v2_consolidated_metadata_empty_dataset( self, memory_store: zarr.storage.MemoryStore ) -> AsyncGroup: - zgroup_bytes = cpu.Buffer.from_bytes(json.dumps({"zarr_format": 2}).encode()) - zmetadata_bytes = cpu.Buffer.from_bytes( - b'{"metadata":{".zgroup":{"zarr_format":2}},"zarr_consolidated_format":1}' + await memory_store.set( + ".zgroup", cpu.Buffer.from_bytes(json.dumps({"zarr_format": 2}).encode()) ) - return AsyncGroup._from_bytes_v2( - StorePath(memory_store, path=""), - zgroup_bytes, - zattrs_bytes=None, - consolidated_metadata_bytes=zmetadata_bytes, + await memory_store.set( + ".zmetadata", + cpu.Buffer.from_bytes( + b'{"metadata":{".zgroup":{"zarr_format":2}},"zarr_consolidated_format":1}' + ), ) + return await AsyncGroup.open(memory_store, zarr_format=2, use_consolidated=True) async def test_consolidated_metadata_backwards_compatibility( self, v2_consolidated_metadata_empty_dataset: AsyncGroup diff --git a/tests/test_open_properties.py b/tests/test_open_properties.py new file mode 100644 index 0000000000..985167c249 --- /dev/null +++ b/tests/test_open_properties.py @@ -0,0 +1,407 @@ +"""The whole option space of `zarr.open`, checked against the invariants its contract implies. + +`_Scenario` names everything that can vary: what is at the path before the call, +and the arguments to `open`. `_expected` is the contract written as an oracle, +independent of the implementation. `check_open` builds the scene, calls `open`, +and checks every invariant. `test_open_option_space` walks every discrete +combination; `test_open_properties` lets Hypothesis vary the parts that are not +discrete (path names, attribute contents) and shrink whatever it finds. +""" + +from __future__ import annotations + +import dataclasses +import itertools +from typing import TYPE_CHECKING, Any, Literal + +import pytest + +import zarr +from zarr import Array, Group +from zarr.abc.store import Store +from zarr.core.common import ZARR_JSON, ZARRAY_JSON, ZATTRS_JSON, ZGROUP_JSON, ZMETADATA_V2_JSON +from zarr.core.sync import sync +from zarr.errors import ( + ArrayNotFoundError, + ContainsArrayError, + ContainsGroupError, + NodeNotFoundError, +) +from zarr.storage import MemoryStore +from zarr.storage._wrapper import WrapperStore + +pytest.importorskip("hypothesis") + +import hypothesis.strategies as st +from hypothesis import event, given + +from zarr.testing.strategies import _attr_keys, _attr_values, node_names + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + from typing import Self + + from zarr.abc.store import ByteRequest + from zarr.core.buffer import Buffer, BufferPrototype + from zarr.core.common import AccessModeLiteral, ZarrFormat + +Existing = Literal["nothing", "array", "group"] +Kind = Literal["array", "group"] + +MODES: tuple[AccessModeLiteral | None, ...] = ("r", "r+", "a", "w", "w-", None) +FORMATS: tuple[ZarrFormat | None, ...] = (2, 3, None) + + +class _RecordingStore(WrapperStore[Store]): + """A store that counts every `get` per key and every write, so a test can see I/O.""" + + get_counts: dict[str, int] + writes: int + + def __init__(self, store: Store) -> None: + super().__init__(store) + self.get_counts = {} + self.writes = 0 + + def _with_store(self, store: Store) -> Self: + # opening in mode "r" makes a read-only copy; it must count into the same tally + new = type(self)(store) + new.get_counts = self.get_counts + new.__dict__["writes_owner"] = self + return new + + def _record_write(self) -> None: + owner = self.__dict__.get("writes_owner", self) + owner.writes += 1 + + def reset(self) -> None: + self.get_counts.clear() + self.writes = 0 + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + self.get_counts[key] = self.get_counts.get(key, 0) + 1 + return await self._store.get(key, prototype, byte_range) + + async def set(self, key: str, value: Buffer) -> None: + self._record_write() + await self._store.set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + self._record_write() + await self._store.set_if_not_exists(key, value) + + async def delete(self, key: str) -> None: + self._record_write() + await self._store.delete(key) + + async def delete_dir(self, prefix: str) -> None: + self._record_write() + await self._store.delete_dir(prefix) + + +@dataclasses.dataclass(frozen=True) +class _Scenario: + existing: Existing + existing_format: ZarrFormat + consolidated: bool + """For a group: whether consolidated metadata was written for it.""" + path: str + mode: AccessModeLiteral | None + zarr_format: ZarrFormat | None + shape: bool + """Whether `shape` (and `dtype`) are passed, which makes the call describe an array.""" + use_consolidated: bool | None + """Only passed without `shape`.""" + attributes: dict[str, Any] + supports_consolidated: bool = True + """Whether the store says it can hold consolidated metadata.""" + + @property + def effective_mode(self) -> AccessModeLiteral: + return self.mode or "a" # the store is writable, so None means "a" + + @property + def visible(self) -> bool: + """Whether the existing node is one the requested format can see.""" + return self.existing != "nothing" and self.zarr_format in (None, self.existing_format) + + @property + def created_format(self) -> ZarrFormat: + return self.zarr_format or 3 + + +# --------------------------------------------------------------------------- +# the contract, as an oracle +# --------------------------------------------------------------------------- + + +def _create_outcome(sc: _Scenario, kind: Kind) -> Kind | type[Exception]: + """What creating `kind` yields once `open` has decided to create. + + Creation checks the path in the format being created, and only that format, + so an existing node of the other format goes unnoticed and the new node is + made alongside it. + """ + if sc.existing != "nothing" and sc.existing_format == sc.created_format: + return ContainsArrayError if sc.existing == "array" else ContainsGroupError + return kind + + +def _expected(sc: _Scenario) -> Kind | type[Exception]: + """The contract of `open`: what comes back, or what is raised.""" + mode = sc.effective_mode + if mode == "w-" and sc.existing != "nothing": + return FileExistsError # anything under the path is enough + if mode in ("w", "w-"): + # "w" empties the path before anything is read; "w-" found it empty + return "array" if sc.shape else "group" + if sc.shape: + # the call describes an array: open_array semantics + if sc.visible and sc.existing == "array": + return "array" + if sc.visible and sc.existing == "group" and sc.existing_format == 3: + return ContainsGroupError # format 3 says what it is in the one document read + # nothing visible as an array: a format 2 group is not looked for + if mode != "a": + return ArrayNotFoundError + return _create_outcome(sc, "array") + if sc.visible: + if sc.existing == "array": + return "array" # whatever was asked about consolidated metadata + if sc.use_consolidated and not sc.supports_consolidated: + return ValueError # asked of a store that can't have it + if sc.use_consolidated and not sc.consolidated: + return ValueError + return "group" + if mode != "a": + return NodeNotFoundError + if sc.use_consolidated and not sc.supports_consolidated: + return ValueError # open_group asks the store before it would create + return _create_outcome(sc, "group") + + +def _opened(sc: _Scenario) -> bool: + """Whether the outcome is the existing node, as opposed to a new one.""" + expected = _expected(sc) + if not isinstance(expected, str): + return False + return sc.visible and sc.effective_mode in ("r", "r+", "a") and (sc.existing == expected) + + +def _expected_reads(sc: _Scenario) -> set[str]: + """The keys `open` reads, each exactly once, when it opens the existing node.""" + if sc.existing_format == 3: + return {ZARR_JSON} + if sc.shape: + keys = {ZARRAY_JSON, ZATTRS_JSON} + else: + keys = {ZARRAY_JSON, ZGROUP_JSON, ZATTRS_JSON} + if sc.use_consolidated is not False and sc.supports_consolidated: + keys.add(ZMETADATA_V2_JSON) + if sc.zarr_format is None: + keys.add(ZARR_JSON) # detecting the format costs the look that finds nothing + return keys + + +# --------------------------------------------------------------------------- +# the check +# --------------------------------------------------------------------------- + + +def _snapshot(store: Store) -> dict[str, bytes]: + from zarr.core.buffer import default_buffer_prototype + + async def _read_all() -> dict[str, bytes]: + out = {} + async for key in store.list(): + buf = await store.get(key, prototype=default_buffer_prototype()) + assert buf is not None + out[key] = buf.to_bytes() + return out + + return sync(_read_all()) + + +class _NoConsolidatedMemoryStore(MemoryStore): + """A store that says it cannot hold consolidated metadata.""" + + @property + def supports_consolidated_metadata(self) -> bool: + return False + + +def _build_scene(sc: _Scenario) -> _RecordingStore: + # the scene is built on a plain store; the store `open` sees may say it cannot + # hold consolidated metadata, over the same contents + inner = MemoryStore() + store = _RecordingStore( + inner if sc.supports_consolidated else _NoConsolidatedMemoryStore(inner._store_dict) + ) + if sc.existing == "array": + zarr.create_array( + inner, + name=sc.path or None, + shape=(3,), + dtype="uint8", + attributes=sc.attributes, + zarr_format=sc.existing_format, + ) + elif sc.existing == "group": + zarr.create_group( + inner, path=sc.path, attributes=sc.attributes, zarr_format=sc.existing_format + ) + if sc.consolidated: + zarr.consolidate_metadata(inner, path=sc.path) + return store + + +def check_open(sc: _Scenario, record: Callable[[str], None] = lambda label: None) -> None: + """Build the scene for `sc`, call `open`, and check every invariant; `record` notes coverage.""" + store = _build_scene(sc) + prefix = f"{sc.path}/" if sc.path else "" + before = _snapshot(store) + store.reset() + kwargs: dict[str, Any] = ( + {"shape": (3,), "dtype": "uint8"} if sc.shape else {"use_consolidated": sc.use_consolidated} + ) + + def call() -> Array[Any] | Group: + return zarr.open( + store=store, mode=sc.mode, zarr_format=sc.zarr_format, path=sc.path, **kwargs + ) + + expected = _expected(sc) + record(f"outcome={expected if isinstance(expected, str) else expected.__name__}") + + if not isinstance(expected, str): + with pytest.raises(expected): + call() + if sc.effective_mode != "w": + # a failed open leaves the store as it was + assert store.writes == 0 + assert _snapshot(store) == before + return + + node = call() + assert isinstance(node, Array if expected == "array" else Group) + assert node.path == sc.path + + if _opened(sc): + record("opened") + assert dict(node.attrs) == sc.attributes + assert node.metadata.zarr_format == sc.existing_format + # only the documents the node needs, each exactly once, and nothing written + assert store.get_counts == {prefix + key: 1 for key in _expected_reads(sc)} + assert store.writes == 0 + assert _snapshot(store) == before + # opening again gives the same node, and so does the opener for that kind + assert call().metadata == node.metadata + same: Array[Any] | Group + if isinstance(node, Array): + same = zarr.open_array( + store=store, mode=sc.effective_mode, zarr_format=sc.zarr_format, path=sc.path + ) + else: + same = zarr.open_group( + store, + mode=sc.effective_mode, + zarr_format=sc.zarr_format, + path=sc.path, + use_consolidated=sc.use_consolidated, + ) + assert (node.metadata.consolidated_metadata is not None) == ( + sc.consolidated and sc.use_consolidated is not False and sc.supports_consolidated + ) + assert same.metadata == node.metadata + else: + record("created") + assert dict(node.attrs) == {} + assert node.metadata.zarr_format == sc.created_format + assert store.writes > 0 + # and it is there to be opened, in the format it was made in + again = zarr.open(store=store, mode="r", zarr_format=sc.created_format, path=sc.path) + assert type(again) is type(node) + assert again.metadata == node.metadata + + +# --------------------------------------------------------------------------- +# every discrete combination +# --------------------------------------------------------------------------- + + +def _discrete_scenarios() -> Iterator[_Scenario]: + existing_nodes: list[tuple[Existing, ZarrFormat, bool]] = [ + ("nothing", 3, False), + *[("array", fmt, False) for fmt in (2, 3)], + *[("group", fmt, cons) for fmt in (2, 3) for cons in (False, True)], + ] + for (existing, fmt, cons), path, mode, zarr_format, supports in itertools.product( + existing_nodes, ("", "outer/inner"), MODES, FORMATS, (True, False) + ): + calls: list[tuple[bool, bool | None]] = [ + (True, None), + *[(False, uc) for uc in (None, False, True)], + ] + for shape, use_consolidated in calls: + yield _Scenario( + existing=existing, + existing_format=fmt, + consolidated=cons, + path=path, + mode=mode, + zarr_format=zarr_format, + shape=shape, + use_consolidated=use_consolidated, + attributes={"old": True}, + supports_consolidated=supports, + ) + + +def _scenario_id(sc: _Scenario) -> str: + node = sc.existing if sc.existing == "nothing" else f"{sc.existing}v{sc.existing_format}" + if sc.consolidated: + node += "c" + call = "shape" if sc.shape else f"uc={sc.use_consolidated}" + store = "" if sc.supports_consolidated else "-nocons" + return f"{node}-{sc.path or 'root'}-mode={sc.mode}-fmt={sc.zarr_format}-{call}{store}" + + +@pytest.mark.parametrize("sc", list(_discrete_scenarios()), ids=_scenario_id) +def test_open_option_space(sc: _Scenario) -> None: + """`open` follows its contract for every combination of what is there and what is asked.""" + check_open(sc) + + +# --------------------------------------------------------------------------- +# and the parts that are not discrete +# --------------------------------------------------------------------------- + +_names = node_names.filter(lambda name: not name.startswith(".")) +_paths = st.just("") | st.lists(_names, min_size=1, max_size=3).map("/".join) +_attributes = st.dictionaries(_attr_keys, _attr_values, max_size=3) + + +@st.composite +def scenarios(draw: st.DrawFn) -> _Scenario: + existing: Existing = draw(st.sampled_from(["nothing", "array", "group"])) + shape = draw(st.booleans()) + return _Scenario( + existing=existing, + existing_format=draw(st.sampled_from([2, 3])), + consolidated=existing == "group" and draw(st.booleans()), + path=draw(_paths), + mode=draw(st.sampled_from(MODES)), + zarr_format=draw(st.sampled_from(FORMATS)), + shape=shape, + use_consolidated=None if shape else draw(st.none() | st.booleans()), + attributes=draw(_attributes), + supports_consolidated=draw(st.booleans()), + ) + + +@given(sc=scenarios()) +def test_open_properties(sc: _Scenario) -> None: + """`open` follows its contract whatever the path and the attributes.""" + check_open(sc, record=event) diff --git a/tests/test_store/test_wrapper.py b/tests/test_store/test_wrapper.py index e556a108c5..bbdd0410c3 100644 --- a/tests/test_store/test_wrapper.py +++ b/tests/test_store/test_wrapper.py @@ -167,3 +167,17 @@ def test_wrapper_delete_sync_without_inner_sync_raises(tmp_path: Any) -> None: store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) with pytest.raises(TypeError, match="does not support synchronous delete"): store.delete_sync("key") + + +def test_wrapper_forwards_supports_consolidated_metadata() -> None: + """A wrapper answers for the store it wraps: one that can't hold consolidated metadata stays that way wrapped.""" + from zarr.storage import MemoryStore + from zarr.storage._wrapper import WrapperStore + + class NoConsolidated(MemoryStore): + @property + def supports_consolidated_metadata(self) -> bool: + return False + + assert WrapperStore(MemoryStore()).supports_consolidated_metadata is True + assert WrapperStore(NoConsolidated()).supports_consolidated_metadata is False