diff --git a/src/zarr/abc/store.py b/src/zarr/abc/store.py index c02dc9d00..e7acce78d 100644 --- a/src/zarr/abc/store.py +++ b/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: """ diff --git a/src/zarr/storage/local.py b/src/zarr/storage/local.py index b8b67003c..f8ff40d52 100644 --- a/src/zarr/storage/local.py +++ b/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,14 +217,14 @@ 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("/") @@ -232,7 +232,7 @@ async def list_prefix(self, prefix: str) -> AsyncGenerator[str]: 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: diff --git a/src/zarr/storage/logging.py b/src/zarr/storage/logging.py index d3e55c068..e79b7fa44 100644 --- a/src/zarr/storage/logging.py +++ b/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): diff --git a/src/zarr/storage/memory.py b/src/zarr/storage/memory.py index df33876df..a62c30ea3 100644 --- a/src/zarr/storage/memory.py +++ b/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("/") diff --git a/src/zarr/storage/remote.py b/src/zarr/storage/remote.py index d471c85f5..962806977 100644 --- a/src/zarr/storage/remote.py +++ b/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 diff --git a/src/zarr/storage/zip.py b/src/zarr/storage/zip.py index a45cc1672..bdae1b597 100644 --- a/src/zarr/storage/zip.py +++ b/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("/")