use common save_metadata function for groups and arrays

Author: K-Meech

src/zarr/core/array.py

@@ -25,7 +25,6 @@
 import zarr
 from zarr.abc.codec import ArrayArrayCodec, ArrayBytesCodec, BytesBytesCodec, Codec
 from zarr.abc.numcodec import Numcodec, _is_numcodec
-from zarr.abc.store import Store, set_or_delete
 from zarr.codecs._v2 import V2Codec
 from zarr.codecs.bytes import BytesCodec
 from zarr.codecs.vlen_utf8 import VLenBytesCodec, VLenUTF8Codec
@@ -109,6 +108,7 @@
     ArrayV3MetadataDict,
     T_ArrayMetadata,
 )
+from zarr.core.metadata.io import save_metadata
 from zarr.core.metadata.v2 import (
     CompressorLikev2,
     get_object_codec_id,
@@ -119,7 +119,6 @@
 from zarr.core.sync import sync
 from zarr.errors import (
     ArrayNotFoundError,
-    ContainsArrayError,
     MetadataValidationError,
     ZarrDeprecationWarning,
     ZarrUserWarning,
@@ -140,9 +139,9 @@
     import numpy.typing as npt
 
     from zarr.abc.codec import CodecPipeline
+    from zarr.abc.store import Store
     from zarr.codecs.sharding import ShardingCodecIndexLocation
     from zarr.core.dtype.wrapper import TBaseDType, TBaseScalar
-    from zarr.core.group import AsyncGroup
     from zarr.storage import StoreLike
 
 
@@ -1496,42 +1495,7 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
         """
         Asynchronously save the array metadata.
         """
-        to_save = metadata.to_buffer_dict(cpu_buffer_prototype)
-        set_awaitables = [
-            set_or_delete(self.store_path / key, value) for key, value in to_save.items()
-        ]
-
-        if ensure_parents:
-            # To enable zarr.create(store, path="a/b/c"), we need to create all the intermediate groups.
-            parents = _build_parents(self)
-            ensure_array_awaitables = []
-
-            for parent in parents:
-                # Error if an array already exists at any parent location. Only groups can have child nodes.
-                ensure_array_awaitables.append(
-                    ensure_no_existing_node(
-                        parent.store_path, metadata.zarr_format, node_type="array"
-                    )
-                )
-                set_awaitables.extend(
-                    [
-                        (parent.store_path / key).set_if_not_exists(value)
-                        for key, value in parent.metadata.to_buffer_dict(
-                            cpu_buffer_prototype
-                        ).items()
-                    ]
-                )
-
-            # Checks for parent arrays must happen first, before any metadata is modified
-            try:
-                await gather(*ensure_array_awaitables)
-            except ContainsArrayError as e:
-                set_awaitables = []  # clear awaitables to avoid printed RuntimeWarning: coroutine was never awaited
-                raise ValueError(
-                    f"A parent of {self.store_path} is an array - only groups may have child nodes."
-                ) from e
-
-        await gather(*set_awaitables)
+        await save_metadata(self.store_path, metadata, ensure_parents=ensure_parents)
 
     async def _set_selection(
         self,
@@ -3899,37 +3863,6 @@ async def chunks_initialized(
     )
 
 
-def _build_parents(
-    node: AsyncArray[ArrayV2Metadata] | AsyncArray[ArrayV3Metadata] | AsyncGroup,
-) -> list[AsyncGroup]:
-    from zarr.core.group import AsyncGroup, GroupMetadata
-
-    store = node.store_path.store
-    path = node.store_path.path
-    if not path:
-        return []
-
-    required_parts = path.split("/")[:-1]
-    parents = [
-        # the root group
-        AsyncGroup(
-            metadata=GroupMetadata(zarr_format=node.metadata.zarr_format),
-            store_path=StorePath(store=store, path=""),
-        )
-    ]
-
-    for i, part in enumerate(required_parts):
-        p = "/".join(required_parts[:i] + [part])
-        parents.append(
-            AsyncGroup(
-                metadata=GroupMetadata(zarr_format=node.metadata.zarr_format),
-                store_path=StorePath(store=store, path=p),
-            )
-        )
-
-    return parents
-
-
 FiltersLike: TypeAlias = (
     Iterable[dict[str, JSON] | ArrayArrayCodec | Numcodec]
     | ArrayArrayCodec

src/zarr/core/group.py

@@ -28,7 +28,6 @@
     FiltersLike,
     SerializerLike,
     ShardsLike,
-    _build_parents,
     _parse_deprecated_compressor,
     create_array,
 )
@@ -49,6 +48,7 @@
 )
 from zarr.core.config import config
 from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata
+from zarr.core.metadata.io import save_metadata
 from zarr.core.sync import SyncMixin, sync
 from zarr.errors import (
     ContainsArrayError,
@@ -808,22 +808,7 @@ async def get(
             return default
 
     async def _save_metadata(self, ensure_parents: bool = False) -> None:
-        to_save = self.metadata.to_buffer_dict(default_buffer_prototype())
-        awaitables = [set_or_delete(self.store_path / key, value) for key, value in to_save.items()]
-
-        if ensure_parents:
-            parents = _build_parents(self)
-            for parent in parents:
-                awaitables.extend(
-                    [
-                        (parent.store_path / key).set_if_not_exists(value)
-                        for key, value in parent.metadata.to_buffer_dict(
-                            default_buffer_prototype()
-                        ).items()
-                    ]
-                )
-
-        await asyncio.gather(*awaitables)
+        await save_metadata(self.store_path, self.metadata, ensure_parents=ensure_parents)
 
     @property
     def path(self) -> str:

src/zarr/core/metadata/io.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING
+
+from zarr.abc.store import set_or_delete
+from zarr.core.buffer.core import default_buffer_prototype
+from zarr.errors import ContainsArrayError
+from zarr.storage._common import StorePath, ensure_no_existing_node
+
+if TYPE_CHECKING:
+    from zarr.core.common import ZarrFormat
+    from zarr.core.group import AsyncGroup, GroupMetadata
+    from zarr.core.metadata import ArrayMetadata
+
+
+def _build_parents(store_path: StorePath, zarr_format: ZarrFormat) -> list[AsyncGroup]:
+    from zarr.core.group import AsyncGroup, GroupMetadata
+
+    store = store_path.store
+    path = store_path.path
+    if not path:
+        return []
+
+    required_parts = path.split("/")[:-1]
+    parents = [
+        # the root group
+        AsyncGroup(
+            metadata=GroupMetadata(zarr_format=zarr_format),
+            store_path=StorePath(store=store, path=""),
+        )
+    ]
+
+    for i, part in enumerate(required_parts):
+        p = "/".join(required_parts[:i] + [part])
+        parents.append(
+            AsyncGroup(
+                metadata=GroupMetadata(zarr_format=zarr_format),
+                store_path=StorePath(store=store, path=p),
+            )
+        )
+
+    return parents
+
+
+async def save_metadata(
+    store_path: StorePath, metadata: ArrayMetadata | GroupMetadata, ensure_parents: bool = False
+) -> None:
+    """Asynchronously save the array or group metadata.
+
+    Parameters
+    ----------
+    store_path : StorePath
+        Location to save metadata.
+    metadata : ArrayMetadata | GroupMetadata
+        Metadata to save.
+    ensure_parents : bool, optional
+        Whether to create any missing parent groups
+
+    Raises
+    ------
+    ValueError
+    """
+    to_save = metadata.to_buffer_dict(default_buffer_prototype())
+    set_awaitables = [set_or_delete(store_path / key, value) for key, value in to_save.items()]
+
+    if ensure_parents:
+        # To enable zarr.create(store, path="a/b/c"), we need to create all the intermediate groups.
+        parents = _build_parents(store_path, metadata.zarr_format)
+        ensure_array_awaitables = []
+
+        for parent in parents:
+            # Error if an array already exists at any parent location. Only groups can have child nodes.
+            ensure_array_awaitables.append(
+                ensure_no_existing_node(parent.store_path, metadata.zarr_format, node_type="array")
+            )
+            set_awaitables.extend(
+                [
+                    (parent.store_path / key).set_if_not_exists(value)
+                    for key, value in parent.metadata.to_buffer_dict(
+                        default_buffer_prototype()
+                    ).items()
+                ]
+            )
+
+        # Checks for parent arrays must happen first, before any metadata is modified
+        try:
+            await asyncio.gather(*ensure_array_awaitables)
+        except ContainsArrayError as e:
+            set_awaitables = []  # clear awaitables to avoid printed RuntimeWarning: coroutine was never awaited
+            raise ValueError(
+                f"A parent of {store_path} is an array - only groups may have child nodes."
+            ) from e
+
+    await asyncio.gather(*set_awaitables)

src/zarr/storage/_common.py

@@ -427,14 +427,21 @@ async def ensure_no_existing_node(
     elif zarr_format == 3:
         extant_node = await _contains_node_v3(store_path)
 
-    if extant_node == "array" and node_type != "group":
-        raise ContainsArrayError(store_path.store, store_path.path)
-    elif extant_node == "group" and node_type != "array":
-        raise ContainsGroupError(store_path.store, store_path.path)
-    elif extant_node == "nothing":
-        return
-    msg = f"Invalid value for extant_node: {extant_node}"
-    raise ValueError(msg)
+    match extant_node:
+        case "array":
+            if node_type != "group":
+                raise ContainsArrayError(store_path.store, store_path.path)
+
+        case "group":
+            if node_type != "array":
+                raise ContainsGroupError(store_path.store, store_path.path)
+
+        case "nothing":
+            return
+
+        case _:
+            msg = f"Invalid value for extant_node: {extant_node}"  # type: ignore[unreachable]
+            raise ValueError(msg)
 
 
 async def _contains_node_v3(store_path: StorePath) -> Literal["array", "group", "nothing"]: