fixed members

Author: TomAugspurger

src/zarr/core/group.py

@@ -1124,7 +1124,16 @@ async def nmembers(
         -------
         count : int
         """
-        if self.metadata.consolidated_metadata is not None:
+        # check if we can use consolidated metadata, which requires that we have non-None
+        # consolidated metadata at all points in the hierarchy.
+        use_consolidated_metadata = self.metadata.consolidated_metadata is not None and all(
+            x.consolidated_metadata is not None
+            for x in self.metadata.consolidated_metadata.flattened_metadata.values()
+            if isinstance(x, GroupMetadata)
+        )
+
+        if use_consolidated_metadata:
+            assert self.metadata.consolidated_metadata is not None  # helping mypy
             return len(self.metadata.consolidated_metadata.flattened_metadata)
         # TODO: consider using aioitertools.builtins.sum for this
         # return await aioitertools.builtins.sum((1 async for _ in self.members()), start=0)
@@ -1174,13 +1183,16 @@ async def _members(
     ]:
         if self.metadata.consolidated_metadata is not None:
             # we should be able to do members without any additional I/O
+            members = self._members_consolidated(max_depth, current_depth)
+
             try:
-                members = self._members_consolidated(max_depth, current_depth)
+                # we already have this in memory, so fine to build this list
+                # and catch the exception if needed.
+                members_ = list(members)
             except _MixedConsolidatedMetadataException:
-                # we've already logged this. We'll fall back to the non-consolidated version.
                 pass
             else:
-                for member in members:
+                for member in members_:
                     yield member
                 return
 
@@ -1238,6 +1250,8 @@ def _members_consolidated(
     ]:
         consolidated_metadata = self.metadata.consolidated_metadata
 
+        if consolidated_metadata is None:
+            raise _MixedConsolidatedMetadataException(prefix)
         # we kind of just want the top-level keys.
         if consolidated_metadata is not None:
             for key in consolidated_metadata.metadata.keys():

tests/v3/test_group.py

@@ -327,9 +327,6 @@ def test_group_getitem(store: Store, zarr_format: ZarrFormat, consolidated: bool
 
     # Now test the mixed case
     if consolidated:
-        group = zarr.api.synchronous.consolidate_metadata(store=store, zarr_format=zarr_format)
-        # we're going to assume that `group.metadata` is correct, and reuse that to focus
-        # on indexing in this test. Other tests verify the correctness of group.metadata
         object.__setattr__(
             group.metadata.consolidated_metadata.metadata["subgroup"],
             "consolidated_metadata",
@@ -1036,6 +1033,20 @@ async def test_group_members_async(store: Store, consolidated_metadata: bool) ->
     with pytest.raises(ValueError, match="max_depth"):
         [x async for x in group.members(max_depth=-1)]
 
+    if consolidated_metadata:
+        # test for mixed
+        object.__setattr__(
+            group.metadata.consolidated_metadata.metadata["g0"].consolidated_metadata.metadata[
+                "g1"
+            ],
+            "consolidated_metadata",
+            None,
+        )
+        all_children = sorted([x async for x in group.members(max_depth=None)], key=lambda x: x[0])
+        assert len(all_children) == len(expected)
+        nmembers = await group.nmembers(max_depth=None)
+        assert nmembers == 6
+
 
 async def test_require_group(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None:
     root = await AsyncGroup.from_store(store=store, zarr_format=zarr_format)