Skip to content

Commit c537fca

Browse files
committed
Factor out code for guessing zarr version
1 parent 24385ab commit c537fca

File tree

1 file changed

+44
-45
lines changed

1 file changed

+44
-45
lines changed

src/zarr/core/group.py

Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,30 @@ def _parse_async_node(
132132
raise TypeError(f"Unknown node type, got {type(node)}")
133133

134134

135+
async def _get_zarr_version(store_path: StorePath) -> Literal[2, 3]:
136+
"""
137+
Guess Zarr format from present metadata files in a Store.
138+
"""
139+
(
140+
zarr_json_bytes,
141+
zgroup_bytes,
142+
) = await asyncio.gather(
143+
(store_path / ZARR_JSON).exists(),
144+
(store_path / ZGROUP_JSON).exists(),
145+
)
146+
if zarr_json_bytes and zgroup_bytes:
147+
# warn and favor v3
148+
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."
149+
warnings.warn(msg, category=ZarrUserWarning, stacklevel=1)
150+
if not zarr_json_bytes and not zgroup_bytes:
151+
raise FileNotFoundError(f"could not find zarr.json or .zgroup objects in {store_path}")
152+
# set zarr_format based on which keys were found
153+
if zarr_json_bytes:
154+
return 3
155+
else:
156+
return 2
157+
158+
135159
@dataclass(frozen=True)
136160
class ConsolidatedMetadata:
137161
"""
@@ -513,6 +537,14 @@ async def open(
513537
to load consolidated metadata from a non-default key.
514538
"""
515539
store_path = await make_store_path(store)
540+
# Guess zarr_format if not passed explicitly
541+
if zarr_format is None:
542+
zarr_format = await _get_zarr_version(store_path)
543+
return await cls.open(
544+
store=store, zarr_format=zarr_format, use_consolidated=use_consolidated
545+
)
546+
assert zarr_format is not None
547+
516548
if not store_path.store.supports_consolidated_metadata:
517549
# Fail if consolidated metadata was requested but the Store doesn't support it
518550
if use_consolidated:
@@ -524,12 +556,11 @@ async def open(
524556
# if use_consolidated was None (optional), the Store dictates it doesn't want consolidation
525557
use_consolidated = False
526558

527-
consolidated_key = ZMETADATA_V2_JSON
528-
529-
if (zarr_format == 2 or zarr_format is None) and isinstance(use_consolidated, str):
530-
consolidated_key = use_consolidated
531-
532559
if zarr_format == 2:
560+
consolidated_key = ZMETADATA_V2_JSON
561+
if isinstance(use_consolidated, str):
562+
consolidated_key = use_consolidated
563+
533564
paths = [store_path / ZGROUP_JSON, store_path / ZATTRS_JSON]
534565
if use_consolidated or use_consolidated is None:
535566
paths.append(store_path / consolidated_key)
@@ -546,43 +577,6 @@ async def open(
546577
else:
547578
maybe_consolidated_metadata_bytes = None
548579

549-
elif zarr_format == 3:
550-
zarr_json_bytes = await (store_path / ZARR_JSON).get()
551-
if zarr_json_bytes is None:
552-
raise FileNotFoundError(store_path)
553-
elif zarr_format is None:
554-
(
555-
zarr_json_bytes,
556-
zgroup_bytes,
557-
zattrs_bytes,
558-
maybe_consolidated_metadata_bytes,
559-
) = await asyncio.gather(
560-
(store_path / ZARR_JSON).get(),
561-
(store_path / ZGROUP_JSON).get(),
562-
(store_path / ZATTRS_JSON).get(),
563-
(store_path / str(consolidated_key)).get(),
564-
)
565-
if zarr_json_bytes is not None and zgroup_bytes is not None:
566-
# warn and favor v3
567-
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."
568-
warnings.warn(msg, category=ZarrUserWarning, stacklevel=1)
569-
if zarr_json_bytes is None and zgroup_bytes is None:
570-
raise FileNotFoundError(
571-
f"could not find zarr.json or .zgroup objects in {store_path}"
572-
)
573-
# set zarr_format based on which keys were found
574-
if zarr_json_bytes is not None:
575-
zarr_format = 3
576-
else:
577-
zarr_format = 2
578-
else:
579-
msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable]
580-
raise MetadataValidationError(msg)
581-
582-
if zarr_format == 2:
583-
# this is checked above, asserting here for mypy
584-
assert zgroup_bytes is not None
585-
586580
if use_consolidated and maybe_consolidated_metadata_bytes is None:
587581
# the user requested consolidated metadata, but it was missing
588582
raise ValueError(consolidated_key)
@@ -595,9 +589,11 @@ async def open(
595589
return cls._from_bytes_v2(
596590
store_path, zgroup_bytes, zattrs_bytes, maybe_consolidated_metadata_bytes
597591
)
598-
else:
599-
# V3 groups are comprised of a zarr.json object
600-
assert zarr_json_bytes is not None
592+
593+
elif zarr_format == 3:
594+
zarr_json_bytes = await (store_path / ZARR_JSON).get()
595+
if zarr_json_bytes is None:
596+
raise FileNotFoundError(store_path)
601597
if not isinstance(use_consolidated, bool | None):
602598
raise TypeError("use_consolidated must be a bool or None for Zarr format 3.")
603599

@@ -606,6 +602,9 @@ async def open(
606602
zarr_json_bytes,
607603
use_consolidated=use_consolidated,
608604
)
605+
else:
606+
msg = f"Invalid value for 'zarr_format'. Expected 2, 3, or None. Got '{zarr_format}'." # type: ignore[unreachable]
607+
raise MetadataValidationError(msg)
609608

610609
@classmethod
611610
def _from_bytes_v2(

0 commit comments

Comments
 (0)