AsyncGenerator[str, None] → AsyncIterator[str]

In the process, make methods of abstract base class async. The rationale
is that the methods of child classes should match the signatures of the
methods of the abstract base class, including async.

The `None` was added in 0e035fb, let's try to remove it.

Author: DimitriPapadopoulos

docs/roadmap.rst

@@ -234,13 +234,13 @@ expose the required methods as async methods.
        async def set_partial_values(self, key_start_values: List[Tuple[str, int, Union[bytes, bytearray, memoryview]]]) -> None:
            ...  # required for writable stores
 
-       async def list(self) -> List[str]:
+       async def list(self) -> AsyncIterator[str]:
            ...  # required for listable stores
 
-       async def list_prefix(self, prefix: str) -> List[str]:
+       async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
            ...  # required for listable stores
 
-       async def list_dir(self, prefix: str) -> List[str]:
+       async def list_dir(self, prefix: str) -> AsyncIterator[str]:
            ...  # required for listable stores
 
        # additional (optional methods)

src/zarr/abc/store.py

@@ -5,7 +5,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,17 +329,16 @@ def supports_listing(self) -> bool:
         ...
 
     @abstractmethod
-    def list(self) -> AsyncGenerator[str, None]:
+    async def list(self) -> AsyncIterator[str]:
         """Retrieve all keys in the store.
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
-        ...
 
     @abstractmethod
-    def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         """
         Retrieve all keys in the store that begin with a given prefix. Keys are returned with the
         common leading prefix removed.
@@ -350,12 +349,11 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
-        ...
 
     @abstractmethod
-    def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+    async 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.
@@ -366,9 +364,8 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
 
         Returns
         -------
-        AsyncGenerator[str, None]
+        AsyncIterator[str]
         """
-        ...
 
     def close(self) -> None:
         """Close the store."""

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,28 +217,27 @@ 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, None]:
+    async def list(self) -> AsyncIterator[str]:
         # docstring inherited
-        to_strip = str(self.root) + "/"
-        for p in list(self.root.rglob("*")):
+        to_strip = f"{self.root}/"
+        for p in await asyncio.to_thread(self.root.rglob, "*"):
             if p.is_file():
                 yield str(p).replace(to_strip, "")
 
-    async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         to_strip = os.path.join(str(self.root / prefix))
-        for p in (self.root / prefix).rglob("*"):
+        for p in await asyncio.to_thread((self.root / prefix).rglob, "*"):
             if p.is_file():
                 yield str(p.relative_to(to_strip))
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         base = self.root / prefix
-        to_strip = str(base) + "/"
+        to_strip = f"{base}/"
 
         try:
-            key_iter = base.iterdir()
-            for key in key_iter:
+            for key in await asyncio.to_thread(base.iterdir):
                 yield str(key).replace(to_strip, "")
         except (FileNotFoundError, NotADirectoryError):
             pass

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, None]:
+    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, None]:
+    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, None]:
+    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

@@ -8,7 +8,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
@@ -143,21 +143,20 @@ async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, by
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str, None]:
+    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, None]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         for key in self._store_dict:
             if key.startswith(prefix):
                 yield key.removeprefix(prefix)
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        if prefix.endswith("/"):
-            prefix = prefix[:-1]
+        prefix = prefix.rstrip("/")
 
         if prefix == "":
             keys_unique = {k.split("/")[0] for k in self._store_dict}

src/zarr/storage/remote.py

@@ -8,7 +8,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
 
@@ -280,13 +280,13 @@ async def set_partial_values(
         # docstring inherited
         raise NotImplementedError
 
-    async def list(self) -> AsyncGenerator[str, None]:
+    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, None]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         prefix = f"{self.path}/{prefix.rstrip('/')}"
         try:
@@ -296,7 +296,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
         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, None]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         find_str = f"{self.path}/{prefix}"
         for onefile in await self.fs._find(find_str, 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"]
 
@@ -231,22 +231,21 @@ async def exists(self, key: str) -> bool:
             else:
                 return True
 
-    async def list(self) -> AsyncGenerator[str, None]:
+    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, None]:
+    async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
         async for key in self.list():
             if key.startswith(prefix):
                 yield key.removeprefix(prefix)
 
-    async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
+    async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        if prefix.endswith("/"):
-            prefix = prefix[:-1]
+        prefix = prefix.rstrip("/")
 
         keys = self._zf.namelist()
         seen = set()