refactor(api): build zarr.open, open_array and open_group on one metadata reader

changes/4366.misc.md

@@ -0,0 +1 @@
+`zarr.open`, `open_array`, `open_group`, `AsyncArray.open` and `AsyncGroup.open` are now built on one metadata reader, `zarr.core.group.read_node_metadata`, composed of a per-format `read_v3_metadata` and `read_v2_metadata`; each reads a node's documents once and returns the parsed metadata whichever kind the node is, and with `zarr_format=None` format 3 is tried first and the format 2 documents are read only when there is no `zarr.json`. Opening a Zarr format 3 node costs a single read where `zarr.open` used to make up to seven. `zarr.open` now has a stated contract: with `shape` it behaves as `open_array`; otherwise the reading modes open whatever node is there ('a' creating a group when there is none) and the creating modes create a group ('w' replacing whatever is there). Consequences: `zarr.open(mode="w")` replaces an existing array with a group instead of returning it; `zarr.open(shape=...)` on a group raises `ContainsGroupError` instead of opening the group; `zarr.open(mode="r")` on nothing raises `NodeNotFoundError`; opening an array as a group raises `ContainsArrayError` and a group as an array `ContainsGroupError` in every format and mode; a Zarr format 3 document whose `node_type` is missing or unknown raises `NodeTypeValidationError` everywhere; and a path holding both formats is read as format 3 without a warning.

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_node,
     create_hierarchy,
 )
-from zarr.core.metadata import ArrayMetadataDict, ArrayV2Metadata
+from zarr.core.metadata import ArrayV2Metadata
 from zarr.errors import (
     ArrayNotFoundError,
+    ContainsArrayError,
+    ContainsGroupError,
     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,23 @@ 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:
+    if "shape" in kwargs:
+        # the call describes an array
         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)
+    if mode in _READ_MODES:
+        node = await _open_node(
+            store_path,
+            zarr_format=zarr_format,
+            use_consolidated=kwargs.get("use_consolidated"),
+            config=kwargs.get("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(
@@ -857,24 +862,21 @@ 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:
+        node = await _open_node(
+            store_path, zarr_format=zarr_format, use_consolidated=use_consolidated
+        )
+        if isinstance(node, AsyncGroup):
+            return node
+        if node is not None:
+            msg = f"An array exists in store {store_path.store} at path {store_path.path}."
+            raise ContainsArrayError(msg)
     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 +1268,26 @@ 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 is here is what the caller gets, unless it is a group. An array
+        # has no consolidated metadata to read.
+        node = await _open_node(
+            store_path, zarr_format=zarr_format, use_consolidated=False, config=kwargs.get("config")
+        )
+        if isinstance(node, AsyncArray):
+            return node
+        if node is not None:
+            msg = f"A group exists in store {store_path.store} at path {store_path.path}."
+            raise ContainsGroupError(msg)
+    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:

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

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,12 @@
     RectilinearChunkGridMetadata,
     RegularChunkGridMetadata,
     create_chunk_grid_metadata,
-    parse_node_type_array,
 )
 from zarr.core.sync import sync
 from zarr.errors import (
     ArrayNotFoundError,
     ChunkNotFoundError,
-    MetadataValidationError,
+    ContainsGroupError,
     ZarrDeprecationWarning,
     ZarrUserWarning,
 )
@@ -158,7 +151,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 +261,36 @@ 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 node there and
+    `ContainsGroupError` if the node is a group.
+    """
+    # group.py imports this module, so the shared reader is imported here
+    from zarr.core.group import GroupMetadata, read_node_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_node_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)
+    if isinstance(metadata, GroupMetadata):
+        msg = f"A group exists in store {store_path.store} at path {store_path.path}."
+        raise ContainsGroupError(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 document at `store_path`, as a dict.
+
+    Prefer `AsyncArray.open`, which builds the array from the parsed metadata
+    directly; this is kept for callers that want the document.
+    """
+    return (await _read_array_metadata(store_path, zarr_format)).to_dict()
 
 
 async def _prepare_overwrite(
@@ -813,10 +772,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: