feature(store): make list_* methods async generators (zarr-developers#110)

* feature(store): make list_* methods async generators

* Update src/zarr/v3/store/memory.py

* Apply suggestions from code review

- simplify code comments
- use `removeprefix` instead of `strip`

---------

Co-authored-by: Davis Bennett <[email protected]>

Author: jhamman

src/zarr/v3/abc/store.py

@@ -1,4 +1,5 @@
 from abc import abstractmethod, ABC
+from collections.abc import AsyncGenerator
 
 from typing import List, Tuple, Optional
 
@@ -106,17 +107,17 @@ def supports_listing(self) -> bool:
         ...
 
     @abstractmethod
-    async def list(self) -> List[str]:
+    async def list(self) -> AsyncGenerator[str, None]:
         """Retrieve all keys in the store.
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
         ...
 
     @abstractmethod
-    async def list_prefix(self, prefix: str) -> List[str]:
+    async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         """Retrieve all keys in the store with a given prefix.
 
         Parameters
@@ -125,12 +126,12 @@ async def list_prefix(self, prefix: str) -> List[str]:
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
         ...
 
     @abstractmethod
-    async def list_dir(self, prefix: str) -> List[str]:
+    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
         Retrieve all keys and prefixes with a given prefix and which do not contain the character
         “/” after the given prefix.
@@ -141,6 +142,6 @@ async def list_dir(self, prefix: str) -> List[str]:
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
         ...

src/zarr/v3/group.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+from collections.abc import AsyncGenerator
 from typing import TYPE_CHECKING
 from dataclasses import asdict, dataclass, field, replace
 
@@ -9,7 +10,6 @@
 if TYPE_CHECKING:
     from typing import (
         Any,
-        AsyncGenerator,
         Literal,
         AsyncIterator,
     )
@@ -157,6 +157,7 @@ async def getitem(
         key: str,
     ) -> AsyncArray | AsyncGroup:
         store_path = self.store_path / key
+        logger.warning("key=%s, store_path=%s", key, store_path)
 
         # if `key` names an object in storage, it cannot be an array or group
         if await store_path.exists():
@@ -302,22 +303,24 @@ async def members(self) -> AsyncGenerator[tuple[str, AsyncArray | AsyncGroup], N
             )
 
             raise ValueError(msg)
-        subkeys = await self.store_path.store.list_dir(self.store_path.path)
-        # would be nice to make these special keys accessible programmatically,
-        # and scoped to specific zarr versions
-        subkeys_filtered = filter(lambda v: v not in ("zarr.json", ".zgroup", ".zattrs"), subkeys)
-        # is there a better way to schedule this?
-        for subkey in subkeys_filtered:
-            try:
-                yield (subkey, await self.getitem(subkey))
-            except KeyError:
-                # keyerror is raised when `subkey` names an object (in the object storage sense),
-                # as opposed to a prefix, in the store under the prefix associated with this group
-                # in which case `subkey` cannot be the name of a sub-array or sub-group.
-                logger.warning(
-                    "Object at %s is not recognized as a component of a Zarr hierarchy.", subkey
-                )
-                pass
+
+        async for key in self.store_path.store.list_dir(self.store_path.path):
+            # these keys are not valid child names so we make sure to skip them
+            # TODO: it would be nice to make these special keys accessible programmatically,
+            # and scoped to specific zarr versions
+            if key not in ("zarr.json", ".zgroup", ".zattrs"):
+                try:
+                    # TODO: performance optimization -- load children concurrently   
+                    child = await self.getitem(key)
+                    yield key, child
+                except KeyError:
+                    # keyerror is raised when `key` names an object (in the object storage sense),
+                    # as opposed to a prefix, in the store under the prefix associated with this group.
+                    # In this case, `key` cannot be the name of a sub-array or sub-group.
+                    logger.warning(
+                        "Object at %s is not recognized as a component of a Zarr hierarchy.", key
+                    )
+                    pass
 
     async def contains(self, member: str) -> bool:
         raise NotImplementedError

src/zarr/v3/store/local.py

@@ -2,6 +2,7 @@
 
 import io
 import shutil
+from collections.abc import AsyncGenerator
 from pathlib import Path
 from typing import Union, Optional, List, Tuple
 
@@ -145,22 +146,19 @@ async def exists(self, key: str) -> bool:
         path = self.root / key
         return await to_thread(path.is_file)
 
-    async def list(self) -> List[str]:
+    async def list(self) -> AsyncGenerator[str, None]:
         """Retrieve all keys in the store.
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
+        for p in self.root.rglob(""):
+            if p.is_file():
+                yield str(p)
 
-        # Q: do we want to return strings or Paths?
-        def _list(root: Path) -> List[str]:
-            files = [str(p) for p in root.rglob("") if p.is_file()]
-            return files
 
-        return await to_thread(_list, self.root)
-
-    async def list_prefix(self, prefix: str) -> List[str]:
+    async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
         """Retrieve all keys in the store with a given prefix.
 
         Parameters
@@ -169,16 +167,14 @@ async def list_prefix(self, prefix: str) -> List[str]:
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
+        for p in (self.root / prefix).rglob("*"):
+            if p.is_file():
+                yield str(p)
 
-        def _list_prefix(root: Path, prefix: str) -> List[str]:
-            files = [str(p) for p in (root / prefix).rglob("*") if p.is_file()]
-            return files
-
-        return await to_thread(_list_prefix, self.root, prefix)
 
-    async def list_dir(self, prefix: str) -> List[str]:
+    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         """
         Retrieve all keys and prefixes with a given prefix and which do not contain the character
         “/” after the given prefix.
@@ -189,15 +185,16 @@ async def list_dir(self, prefix: str) -> List[str]:
 
         Returns
         -------
-        list[str]
+        AsyncGenerator[str, None]
         """
+        base = self.root / prefix
+        to_strip = str(base) + "/"
+        
+        try:
+            key_iter = base.iterdir()
+        except (FileNotFoundError, NotADirectoryError):
+            key_iter = []
 
-        def _list_dir(root: Path, prefix: str) -> List[str]:
-            base = root / prefix
-            to_strip = str(base) + "/"
-            try:
-                return [str(key).replace(to_strip, "") for key in base.iterdir()]
-            except (FileNotFoundError, NotADirectoryError):
-                return []
+        for key in key_iter:
+            yield str(key).replace(to_strip, "")
 
-        return await to_thread(_list_dir, self.root, prefix)

src/zarr/v3/store/memory.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import AsyncGenerator
 from typing import Optional, MutableMapping, List, Tuple
 
 from zarr.v3.common import BytesLike
@@ -67,20 +68,16 @@ async def delete(self, key: str) -> None:
     async def set_partial_values(self, key_start_values: List[Tuple[str, int, bytes]]) -> None:
         raise NotImplementedError
 
-    async def list(self) -> List[str]:
-        return list(self._store_dict.keys())
+    async def list(self) -> AsyncGenerator[str, None]:
+        for key in self._store_dict:
+            yield key
 
-    async def list_prefix(self, prefix: str) -> List[str]:
-        return [key for key in self._store_dict if key.startswith(prefix)]
+    async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+        for key in self._store_dict:
+            if key.startswith(prefix):
+                yield key
 
-    async def list_dir(self, prefix: str) -> List[str]:
-        if prefix == "":
-            return list({key.split("/", maxsplit=1)[0] for key in self._store_dict})
-        else:
-            return list(
-                {
-                    key.removeprefix(prefix + "/").split("/")[0]
-                    for key in self._store_dict
-                    if (key.startswith(prefix + "/") and key != prefix)
-                }
-            )
+    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+        for key in self._store_dict:
+            if key.startswith(prefix + "/") and key != prefix:
+                yield key.removeprefix(prefix + "/").rsplit("/", maxsplit=1)[0]