Fixed MemoryStore.list_dir

src/zarr/store/memory.py

@@ -114,9 +114,15 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
 
         if prefix == "":
             keys_unique = set(k.split("/")[0] for k in self._store_dict.keys())
-            for key in keys_unique:
-                yield key
         else:
-            for key in self._store_dict:
-                if key.startswith(prefix + "/") and key != prefix:
-                    yield key.removeprefix(prefix + "/").split("/")[0]
+            # Our dictionary doesn't contain directory markers, but we want to include
+            # a pseudo directory when there's a nested item and we're listing an
+            # intermediate level.
+            keys_unique = {
+                key.removeprefix(prefix + "/").split("/")[0]
+                for key in self._store_dict
+                if key.startswith(prefix + "/") and key != prefix
+            }
+
+        for key in keys_unique:
+            yield key

src/zarr/store/remote.py

@@ -202,7 +202,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         except FileNotFoundError:
             return
         for onefile in (a.replace(prefix + "/", "") for a in allfiles):
-            yield onefile
+            yield onefile.removeprefix(self.path).removeprefix("/")
 
     async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         for onefile in await self._fs._ls(prefix, detail=False):

src/zarr/testing/store.py

@@ -191,9 +191,17 @@ async def test_list_dir(self, store: S) -> None:
         assert out == []
         assert [k async for k in store.list_dir("foo")] == []
         await store.set("foo/zarr.json", Buffer.from_bytes(b"bar"))
-        await store.set("foo/c/1", Buffer.from_bytes(b"\x01"))
+        await store.set("group-0/zarr.json", Buffer.from_bytes(b"\x01"))  # group
+        await store.set("group-0/group-1/zarr.json", Buffer.from_bytes(b"\x01"))  # group
+        await store.set("group-0/group-1/a1/zarr.json", Buffer.from_bytes(b"\x01"))
+        await store.set("group-0/group-1/a2/zarr.json", Buffer.from_bytes(b"\x01"))
+        await store.set("group-0/group-1/a3/zarr.json", Buffer.from_bytes(b"\x01"))
 
-        keys_expected = ["zarr.json", "c"]
+        keys_expected = ["foo", "group-0"]
+        keys_observed = [k async for k in store.list_dir("")]
+        assert set(keys_observed) == set(keys_expected)
+
+        keys_expected = ["zarr.json"]
         keys_observed = [k async for k in store.list_dir("foo")]
 
         assert len(keys_observed) == len(keys_expected), keys_observed
@@ -202,3 +210,23 @@ async def test_list_dir(self, store: S) -> None:
         keys_observed = [k async for k in store.list_dir("foo/")]
         assert len(keys_expected) == len(keys_observed), keys_observed
         assert set(keys_observed) == set(keys_expected), keys_observed
+
+        keys_observed = [k async for k in store.list_dir("group-0")]
+        keys_expected = ["zarr.json", "group-1"]
+
+        assert len(keys_observed) == len(keys_expected), keys_observed
+        assert set(keys_observed) == set(keys_expected), keys_observed
+
+        keys_observed = [k async for k in store.list_dir("group-0/")]
+        assert len(keys_expected) == len(keys_observed), keys_observed
+        assert set(keys_observed) == set(keys_expected), keys_observed
+
+        keys_observed = [k async for k in store.list_dir("group-0/group-1")]
+        keys_expected = ["zarr.json", "a1", "a2", "a3"]
+
+        assert len(keys_observed) == len(keys_expected), keys_observed
+        assert set(keys_observed) == set(keys_expected), keys_observed
+
+        keys_observed = [k async for k in store.list_dir("group-0/group-1")]
+        assert len(keys_expected) == len(keys_observed), keys_observed
+        assert set(keys_observed) == set(keys_expected), keys_observed