Add

Author: dgegen

src/zarr/core/array.py

@@ -196,42 +196,28 @@ def create_codec_pipeline(metadata: ArrayMetadata) -> CodecPipeline:
     raise TypeError  # pragma: no cover
 
 
-def is_concurrency_safe(store: Store) -> bool:
-    fs = getattr(store, "fs", None)
-    return getattr(fs, "asynchronous", True)
-
-
 async def get_array_metadata(
     store_path: StorePath, zarr_format: ZarrFormat | None = 3
 ) -> dict[str, JSON]:
-    concurrency_safe = is_concurrency_safe(store_path.store)
-
     if zarr_format == 2:
-        if concurrency_safe:
-            zarray_bytes, zattrs_bytes = await gather(
-                (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype),
-                (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype),
-            )
-        else:
-            zarray_bytes = await (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype)
-            zattrs_bytes = await (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype)
+        zarray_bytes, zattrs_bytes = await store_path.get_many(
+            ZARRAY_JSON,
+            ZATTRS_JSON,
+            prototype=cpu_buffer_prototype,
+        )
         if zarray_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format == 3:
         zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype)
         if zarr_json_bytes is None:
             raise FileNotFoundError(store_path)
     elif zarr_format is None:
-        if concurrency_safe:
-            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),
-            )
-        else:
-            zarr_json_bytes = await (store_path / ZARR_JSON).get(prototype=cpu_buffer_prototype)
-            zarray_bytes = await (store_path / ZARRAY_JSON).get(prototype=cpu_buffer_prototype)
-            zattrs_bytes = await (store_path / ZATTRS_JSON).get(prototype=cpu_buffer_prototype)
+        zarr_json_bytes, zarray_bytes, zattrs_bytes = await store_path.get_many(
+            ZARR_JSON,
+            ZARRAY_JSON,
+            ZATTRS_JSON,
+            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."
@@ -1445,12 +1431,7 @@ async def _save_metadata(self, metadata: ArrayMetadata, ensure_parents: bool = F
                         ).items()
                     ]
                 )
-
-        if is_concurrency_safe(self.store_path.store):
-            await gather(*awaitables)
-        else:
-            for awaitable in awaitables:
-                await awaitable
+        await gather(*awaitables)
 
     async def _set_selection(
         self,

src/zarr/core/group.py

@@ -31,7 +31,6 @@
     _build_parents,
     _parse_deprecated_compressor,
     create_array,
-    is_concurrency_safe,
 )
 from zarr.core.attributes import Attributes
 from zarr.core.buffer import default_buffer_prototype
@@ -536,23 +535,17 @@ async def open(
             if zarr_json_bytes is None:
                 raise FileNotFoundError(store_path)
         elif zarr_format is None:
-            paths = [
-                (store_path / ZARR_JSON).get(),
-                (store_path / ZGROUP_JSON).get(),
-                (store_path / ZATTRS_JSON).get(),
-                (store_path / consolidated_key).get(),
-            ]
             (
                 zarr_json_bytes,
                 zgroup_bytes,
                 zattrs_bytes,
                 maybe_consolidated_metadata_bytes,
-            ) = (
-                await asyncio.gather(*paths)
-                if is_concurrency_safe(store_path.store)
-                else [await path for path in paths]
+            ) = await store_path.get_many(
+                ZARR_JSON,
+                ZGROUP_JSON,
+                ZATTRS_JSON,
+                consolidated_key,
             )
-
             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."
@@ -3484,23 +3477,11 @@ async def _read_metadata_v2(store: Store, path: str) -> ArrayV2Metadata | GroupM
     # TODO: consider first fetching array metadata, and only fetching group metadata when we don't
     # find an array
     print(f"Reading metadata from {path} in store {store}", file=sys.stderr)
-    if is_concurrency_safe(store):
-        zarray_bytes, zgroup_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()),
-        )
-    else:
-        zarray_bytes = await store.get(
-            _join_paths([path, ZARRAY_JSON]), prototype=default_buffer_prototype()
-        )
-        zgroup_bytes = await store.get(
-            _join_paths([path, ZGROUP_JSON]), prototype=default_buffer_prototype()
-        )
-        zattrs_bytes = await store.get(
-            _join_paths([path, ZATTRS_JSON]), prototype=default_buffer_prototype()
-        )
-
+    zarray_bytes, zgroup_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()),
+    )
     if zattrs_bytes is None:
         zattrs = {}
     else:

src/zarr/storage/_common.py

@@ -2,6 +2,7 @@
 
 import importlib.util
 import json
+from asyncio import gather
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, Self, TypeAlias
 
@@ -163,6 +164,24 @@ async def get(
             prototype = default_buffer_prototype()
         return await self.store.get(self.path, prototype=prototype, byte_range=byte_range)
 
+    async def get_many(
+        self,
+        *suffixes : str,
+        prototype: BufferPrototype | None = None,
+        byte_range: ByteRequest | None = None,
+    ):
+        tasks = [
+            (self / suffix).get(prototype=prototype, byte_range=byte_range) for suffix in suffixes
+        ]
+        if await self._is_concurrency_save():
+            return await gather(*tasks)
+        else:
+            results = []
+            for task in tasks:
+                result = await task
+                results.append(result)
+            return results
+
     async def set(self, value: Buffer, byte_range: ByteRequest | None = None) -> None:
         """
         Write bytes to the store.
@@ -263,6 +282,10 @@ def __eq__(self, other: object) -> bool:
             pass
         return False
 
+    async def _is_concurrency_save(self):
+        fs = getattr(self.store, "fs", None)
+        return getattr(fs, "asynchronous", True)
+
 
 StoreLike: TypeAlias = Store | StorePath | FSMap | Path | str | dict[str, Buffer]