AsyncGenerator → AsyncIterator (zarr-developers#2381)

Author: DimitriPapadopoulos

src/zarr/abc/store.py

@@ -6,7 +6,7 @@
 from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncGenerator, AsyncIterator, Iterable
     from types import TracebackType
     from typing import Any, Self, TypeAlias
 
@@ -329,16 +329,19 @@ def supports_listing(self) -> bool:
         ...
 
     @abstractmethod
-    def list(self) -> AsyncGenerator[str]:
+    def list(self) -> AsyncIterator[str]:
         """Retrieve all keys in the store.
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     @abstractmethod
-    def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         """
         Retrieve all keys in the store that begin with a given prefix. Keys are returned relative
         to the root of the store.
@@ -349,11 +352,14 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     @abstractmethod
-    def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    def list_dir(self, prefix: str) -> AsyncIterator[str]:
         """
         Retrieve all keys and prefixes with a given prefix and which do not contain the character
         “/” after the given prefix.
@@ -364,8 +370,11 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
+        # This method should be async, like overridden methods in child classes.
+        # However, that's not straightforward:
+        # https://stackoverflow.com/questions/68905848
 
     async def delete_dir(self, prefix: str) -> None:
         """

src/zarr/storage/local.py

@@ -12,7 +12,7 @@
 from zarr.core.common import concurrent_map
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncIterator, Iterable
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import AccessModeLiteral
@@ -217,22 +217,22 @@ async def exists(self, key: str) -> bool:
         path = self.root / key
         return await asyncio.to_thread(path.is_file)
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         to_strip = self.root.as_posix() + "/"
         for p in list(self.root.rglob("*")):
             if p.is_file():
                 yield p.as_posix().replace(to_strip, "")
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         to_strip = self.root.as_posix() + "/"
         prefix = prefix.rstrip("/")
         for p in (self.root / prefix).rglob("*"):
             if p.is_file():
                 yield p.as_posix().replace(to_strip, "")
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         base = self.root / prefix
         try:

src/zarr/storage/logging.py

@@ -10,7 +10,7 @@
 from zarr.abc.store import AccessMode, ByteRangeRequest, Store
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Generator, Iterable
+    from collections.abc import AsyncIterator, Generator, Iterable
 
     from zarr.core.buffer import Buffer, BufferPrototype
     from zarr.core.common import AccessModeLiteral
@@ -204,19 +204,19 @@ async def set_partial_values(
         with self.log(keys):
             return await self._store.set_partial_values(key_start_values=key_start_values)
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         with self.log():
             async for key in self._store.list():
                 yield key
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         with self.log(prefix):
             async for key in self._store.list_prefix(prefix=prefix):
                 yield key
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         with self.log(prefix):
             async for key in self._store.list_dir(prefix=prefix):

src/zarr/storage/memory.py

@@ -9,7 +9,7 @@
 from zarr.storage._utils import _normalize_interval_index
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable, MutableMapping
+    from collections.abc import AsyncIterator, Iterable, MutableMapping
 
     from zarr.core.buffer import BufferPrototype
     from zarr.core.common import AccessModeLiteral
@@ -147,19 +147,19 @@ async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, by
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         for key in self._store_dict:
             yield key
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix)
         for key in list(self._store_dict):
             if key.startswith(prefix):
                 yield key
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = prefix.rstrip("/")
 

src/zarr/storage/remote.py

@@ -7,7 +7,7 @@
 from zarr.storage.common import _dereference_path
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncIterator, Iterable
 
     from fsspec.asyn import AsyncFileSystem
 
@@ -322,13 +322,13 @@ async def set_partial_values(
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         allfiles = await self.fs._find(self.path, detail=False, withdirs=False)
         for onefile in (a.replace(self.path + "/", "") for a in allfiles):
             yield onefile
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = f"{self.path}/{prefix.rstrip('/')}"
         try:
@@ -338,7 +338,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
         for onefile in (a.replace(prefix + "/", "") for a in allfiles):
             yield onefile.removeprefix(self.path).removeprefix("/")
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         for onefile in await self.fs._find(
             f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False

src/zarr/storage/zip.py

@@ -11,7 +11,7 @@
 from zarr.core.buffer import Buffer, BufferPrototype
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterable
+    from collections.abc import AsyncIterator, Iterable
 
 ZipStoreAccessModeLiteral = Literal["r", "w", "a"]
 
@@ -234,19 +234,19 @@ async def exists(self, key: str) -> bool:
             else:
                 return True
 
-    async def list(self) -> AsyncGenerator[str]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
         with self._lock:
             for key in self._zf.namelist():
                 yield key
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         async for key in self.list():
             if key.startswith(prefix):
                 yield key
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = prefix.rstrip("/")