Refactor AsyncGroup.open()

src/zarr/core/group.py

@@ -487,7 +487,8 @@ async def open(
         zarr_format: ZarrFormat | None = 3,
         use_consolidated: bool | str | None = None,
     ) -> AsyncGroup:
-        """Open a new AsyncGroup
+        """
+        Create a new AsyncGroup from an existing group.
 
         Parameters
         ----------
@@ -512,6 +513,34 @@ async def open(
             to load consolidated metadata from a non-default key.
         """
         store_path = await make_store_path(store)
+
+        zarr_json_bytes = None
+        zarr_group_bytes = None
+
+        # Guess zarr_format if not passed explicitly
+        if zarr_format is None:
+            (
+                zarr_json_bytes,
+                zarr_group_bytes,
+            ) = await asyncio.gather(
+                (store_path / ZARR_JSON).get(),
+                (store_path / ZGROUP_JSON).get(),
+            )
+            # set zarr_format based on which keys were found
+            if zarr_json_bytes is not None:
+                zarr_format = 3
+                if zarr_group_bytes is not None:
+                    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)
+            elif zarr_group_bytes is not None:
+                zarr_format = 2
+            else:
+                raise FileNotFoundError(
+                    f"could not find zarr.json or .zgroup objects in {store_path}"
+                )
+
+        assert zarr_format is not None
+
         if not store_path.store.supports_consolidated_metadata:
             # Fail if consolidated metadata was requested but the Store doesn't support it
             if use_consolidated:
@@ -523,80 +552,43 @@ async def open(
             # 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)
+            consolidated_key = ZMETADATA_V2_JSON
+            if isinstance(use_consolidated, str):
+                consolidated_key = use_consolidated
 
-            zgroup_bytes, zattrs_bytes, *rest = await asyncio.gather(
-                *[path.get() for path in paths]
-            )
-            if zgroup_bytes is None:
+            if zarr_group_bytes is None:
+                zarr_group_bytes = await (store_path / ZGROUP_JSON).get()
+            if zarr_group_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:
+            zattrs_bytes = await (store_path / ZATTRS_JSON).get()
+            if zattrs_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:
-            # this is checked above, asserting here for mypy
-            assert zgroup_bytes is not None
+            if use_consolidated or use_consolidated is None:
+                consolidated_metadata_bytes = await (store_path / consolidated_key).get()
+            else:
+                consolidated_metadata_bytes = None
 
-            if use_consolidated and maybe_consolidated_metadata_bytes is None:
+            if use_consolidated and 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
+                consolidated_metadata_bytes = None
 
             return cls._from_bytes_v2(
-                store_path, zgroup_bytes, zattrs_bytes, maybe_consolidated_metadata_bytes
+                store_path, zarr_group_bytes, zattrs_bytes, consolidated_metadata_bytes
             )
-        else:
-            # V3 groups are comprised of a zarr.json object
-            assert zarr_json_bytes is not None
+
+        elif zarr_format == 3:
+            if zarr_json_bytes is None:
+                zarr_json_bytes = await (store_path / ZARR_JSON).get()
+            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.")
 
@@ -605,6 +597,9 @@ async def open(
                 zarr_json_bytes,
                 use_consolidated=use_consolidated,
             )
+        else:
+            msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'."  # type: ignore[unreachable]
+            raise MetadataValidationError(msg)
 
     @classmethod
     def _from_bytes_v2(