Skip to content

Commit 97d384a

Browse files
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.
1 parent 5f3a512 commit 97d384a

File tree

7 files changed

+37
-43
lines changed

7 files changed

+37
-43
lines changed

docs/roadmap.rst

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -234,13 +234,13 @@ expose the required methods as async methods.
234234
async def set_partial_values(self, key_start_values: List[Tuple[str, int, Union[bytes, bytearray, memoryview]]]) -> None:
235235
... # required for writable stores
236236
237-
async def list(self) -> List[str]:
237+
async def list(self) -> AsyncIterator[str]:
238238
... # required for listable stores
239239
240-
async def list_prefix(self, prefix: str) -> List[str]:
240+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
241241
... # required for listable stores
242242
243-
async def list_dir(self, prefix: str) -> List[str]:
243+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
244244
... # required for listable stores
245245
246246
# additional (optional methods)

src/zarr/abc/store.py

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import TYPE_CHECKING, NamedTuple, Protocol, runtime_checkable
66

77
if TYPE_CHECKING:
8-
from collections.abc import AsyncGenerator, Iterable
8+
from collections.abc import AsyncGenerator, AsyncIterator, Iterable
99
from types import TracebackType
1010
from typing import Any, Self, TypeAlias
1111

@@ -329,17 +329,16 @@ def supports_listing(self) -> bool:
329329
...
330330

331331
@abstractmethod
332-
def list(self) -> AsyncGenerator[str, None]:
332+
async def list(self) -> AsyncIterator[str]:
333333
"""Retrieve all keys in the store.
334334
335335
Returns
336336
-------
337-
AsyncGenerator[str, None]
337+
AsyncIterator[str]
338338
"""
339-
...
340339

341340
@abstractmethod
342-
def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
341+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
343342
"""
344343
Retrieve all keys in the store that begin with a given prefix. Keys are returned with the
345344
common leading prefix removed.
@@ -350,12 +349,11 @@ def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
350349
351350
Returns
352351
-------
353-
AsyncGenerator[str, None]
352+
AsyncIterator[str]
354353
"""
355-
...
356354

357355
@abstractmethod
358-
def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
356+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
359357
"""
360358
Retrieve all keys and prefixes with a given prefix and which do not contain the character
361359
“/” after the given prefix.
@@ -366,9 +364,8 @@ def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
366364
367365
Returns
368366
-------
369-
AsyncGenerator[str, None]
367+
AsyncIterator[str]
370368
"""
371-
...
372369

373370
def close(self) -> None:
374371
"""Close the store."""

src/zarr/storage/local.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
from zarr.core.common import concurrent_map
1313

1414
if TYPE_CHECKING:
15-
from collections.abc import AsyncGenerator, Iterable
15+
from collections.abc import AsyncIterator, Iterable
1616

1717
from zarr.core.buffer import BufferPrototype
1818
from zarr.core.common import AccessModeLiteral
@@ -217,28 +217,27 @@ async def exists(self, key: str) -> bool:
217217
path = self.root / key
218218
return await asyncio.to_thread(path.is_file)
219219

220-
async def list(self) -> AsyncGenerator[str, None]:
220+
async def list(self) -> AsyncIterator[str]:
221221
# docstring inherited
222-
to_strip = str(self.root) + "/"
223-
for p in list(self.root.rglob("*")):
222+
to_strip = f"{self.root}/"
223+
for p in await asyncio.to_thread(self.root.rglob, "*"):
224224
if p.is_file():
225225
yield str(p).replace(to_strip, "")
226226

227-
async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
227+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
228228
# docstring inherited
229229
to_strip = os.path.join(str(self.root / prefix))
230-
for p in (self.root / prefix).rglob("*"):
230+
for p in await asyncio.to_thread((self.root / prefix).rglob, "*"):
231231
if p.is_file():
232232
yield str(p.relative_to(to_strip))
233233

234-
async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
234+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
235235
# docstring inherited
236236
base = self.root / prefix
237-
to_strip = str(base) + "/"
237+
to_strip = f"{base}/"
238238

239239
try:
240-
key_iter = base.iterdir()
241-
for key in key_iter:
240+
for key in await asyncio.to_thread(base.iterdir):
242241
yield str(key).replace(to_strip, "")
243242
except (FileNotFoundError, NotADirectoryError):
244243
pass

src/zarr/storage/logging.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from zarr.abc.store import AccessMode, ByteRangeRequest, Store
1111

1212
if TYPE_CHECKING:
13-
from collections.abc import AsyncGenerator, Generator, Iterable
13+
from collections.abc import AsyncIterator, Generator, Iterable
1414

1515
from zarr.core.buffer import Buffer, BufferPrototype
1616
from zarr.core.common import AccessModeLiteral
@@ -204,19 +204,19 @@ async def set_partial_values(
204204
with self.log(keys):
205205
return await self._store.set_partial_values(key_start_values=key_start_values)
206206

207-
async def list(self) -> AsyncGenerator[str, None]:
207+
async def list(self) -> AsyncIterator[str]:
208208
# docstring inherited
209209
with self.log():
210210
async for key in self._store.list():
211211
yield key
212212

213-
async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
213+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
214214
# docstring inherited
215215
with self.log(prefix):
216216
async for key in self._store.list_prefix(prefix=prefix):
217217
yield key
218218

219-
async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
219+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
220220
# docstring inherited
221221
with self.log(prefix):
222222
async for key in self._store.list_dir(prefix=prefix):

src/zarr/storage/memory.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from zarr.storage._utils import _normalize_interval_index
99

1010
if TYPE_CHECKING:
11-
from collections.abc import AsyncGenerator, Iterable, MutableMapping
11+
from collections.abc import AsyncIterator, Iterable, MutableMapping
1212

1313
from zarr.core.buffer import BufferPrototype
1414
from zarr.core.common import AccessModeLiteral
@@ -143,21 +143,20 @@ async def set_partial_values(self, key_start_values: Iterable[tuple[str, int, by
143143
# docstring inherited
144144
raise NotImplementedError
145145

146-
async def list(self) -> AsyncGenerator[str, None]:
146+
async def list(self) -> AsyncIterator[str]:
147147
# docstring inherited
148148
for key in self._store_dict:
149149
yield key
150150

151-
async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
151+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
152152
# docstring inherited
153153
for key in self._store_dict:
154154
if key.startswith(prefix):
155155
yield key.removeprefix(prefix)
156156

157-
async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
157+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
158158
# docstring inherited
159-
if prefix.endswith("/"):
160-
prefix = prefix[:-1]
159+
prefix = prefix.rstrip("/")
161160

162161
if prefix == "":
163162
keys_unique = {k.split("/")[0] for k in self._store_dict}

src/zarr/storage/remote.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from zarr.storage.common import _dereference_path
99

1010
if TYPE_CHECKING:
11-
from collections.abc import AsyncGenerator, Iterable
11+
from collections.abc import AsyncIterator, Iterable
1212

1313
from fsspec.asyn import AsyncFileSystem
1414

@@ -280,13 +280,13 @@ async def set_partial_values(
280280
# docstring inherited
281281
raise NotImplementedError
282282

283-
async def list(self) -> AsyncGenerator[str, None]:
283+
async def list(self) -> AsyncIterator[str]:
284284
# docstring inherited
285285
allfiles = await self.fs._find(self.path, detail=False, withdirs=False)
286286
for onefile in (a.replace(self.path + "/", "") for a in allfiles):
287287
yield onefile
288288

289-
async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
289+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
290290
# docstring inherited
291291
prefix = f"{self.path}/{prefix.rstrip('/')}"
292292
try:
@@ -296,7 +296,7 @@ async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
296296
for onefile in (a.replace(prefix + "/", "") for a in allfiles):
297297
yield onefile.removeprefix(self.path).removeprefix("/")
298298

299-
async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
299+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
300300
# docstring inherited
301301
find_str = f"{self.path}/{prefix}"
302302
for onefile in await self.fs._find(find_str, detail=False, maxdepth=None, withdirs=False):

src/zarr/storage/zip.py

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from zarr.core.buffer import Buffer, BufferPrototype
1212

1313
if TYPE_CHECKING:
14-
from collections.abc import AsyncGenerator, Iterable
14+
from collections.abc import AsyncIterator, Iterable
1515

1616
ZipStoreAccessModeLiteral = Literal["r", "w", "a"]
1717

@@ -231,22 +231,21 @@ async def exists(self, key: str) -> bool:
231231
else:
232232
return True
233233

234-
async def list(self) -> AsyncGenerator[str, None]:
234+
async def list(self) -> AsyncIterator[str]:
235235
# docstring inherited
236236
with self._lock:
237237
for key in self._zf.namelist():
238238
yield key
239239

240-
async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]:
240+
async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
241241
# docstring inherited
242242
async for key in self.list():
243243
if key.startswith(prefix):
244244
yield key.removeprefix(prefix)
245245

246-
async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]:
246+
async def list_dir(self, prefix: str) -> AsyncIterator[str]:
247247
# docstring inherited
248-
if prefix.endswith("/"):
249-
prefix = prefix[:-1]
248+
prefix = prefix.rstrip("/")
250249

251250
keys = self._zf.namelist()
252251
seen = set()

0 commit comments

Comments
 (0)