Fix incorrect Group.nmembers for consolidated metadata (#3287)

Author: csparker247

changes/3287.bugfix.rst

@@ -0,0 +1 @@
+Fixes Group.nmembers() ignoring depth when using consolidated metadata.

src/zarr/core/group.py

@@ -1307,7 +1307,18 @@ async def nmembers(
         # check if we can use consolidated metadata, which requires that we have non-None
         # consolidated metadata at all points in the hierarchy.
         if self.metadata.consolidated_metadata is not None:
-            return len(self.metadata.consolidated_metadata.flattened_metadata)
+            if max_depth is not None and max_depth < 0:
+                raise ValueError(f"max_depth must be None or >= 0. Got '{max_depth}' instead")
+            if max_depth is None:
+                return len(self.metadata.consolidated_metadata.flattened_metadata)
+            else:
+                return len(
+                    [
+                        x
+                        for x in self.metadata.consolidated_metadata.flattened_metadata
+                        if x.count("/") <= max_depth
+                    ]
+                )
         # TODO: consider using aioitertools.builtins.sum for this
         # return await aioitertools.builtins.sum((1 async for _ in self.members()), start=0)
         n = 0

tests/test_group.py

@@ -1118,12 +1118,22 @@ async def test_group_members_async(store: Store, consolidated_metadata: bool) ->
             "consolidated_metadata",
             None,
         )
+        # test depth=0
+        nmembers = await group.nmembers(max_depth=0)
+        assert nmembers == 2
+        # test depth=1
+        nmembers = await group.nmembers(max_depth=1)
+        assert nmembers == 4
+        # test depth=None
         all_children = sorted(
             [x async for x in group.members(max_depth=None)], key=operator.itemgetter(0)
         )
         assert len(all_children) == 4
         nmembers = await group.nmembers(max_depth=None)
         assert nmembers == 4
+        # test depth<0
+        with pytest.raises(ValueError, match="max_depth"):
+            await group.nmembers(max_depth=-1)
 
 
 async def test_require_group(store: LocalStore | MemoryStore, zarr_format: ZarrFormat) -> None: