Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

AsyncGenerator → AsyncIterator #2381

Merged
merged 1 commit into from
Nov 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 16 additions & 7 deletions src/zarr/abc/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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:
"""
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/storage/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/storage/logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/storage/memory.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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("/")

Expand Down
8 changes: 4 additions & 4 deletions src/zarr/storage/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
8 changes: 4 additions & 4 deletions src/zarr/storage/zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down Expand Up @@ -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("/")

Expand Down