diff --git a/changes/3287.bugfix.rst b/changes/3287.bugfix.rst new file mode 100644 index 0000000000..a4eaa35312 --- /dev/null +++ b/changes/3287.bugfix.rst @@ -0,0 +1 @@ +Fixes Group.nmembers() ignoring depth when using consolidated metadata. diff --git a/src/zarr/core/group.py b/src/zarr/core/group.py index a868ee31fa..a398aa01aa 100644 --- a/src/zarr/core/group.py +++ b/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 diff --git a/tests/test_group.py b/tests/test_group.py index ac1afb539b..7705fa205a 100644 --- a/tests/test_group.py +++ b/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: