From 1e5d05c2d25e0bcb2fe83b73e72892bbdf35dd0f Mon Sep 17 00:00:00 2001 From: "Michael V. DePalatis" Date: Tue, 30 Oct 2018 19:50:57 -0600 Subject: [PATCH 1/4] Implement as_completed as an asynchronous iterator --- Lib/asyncio/tasks.py | 83 ++++++++++++++++++++++---------------------- 1 file changed, 41 insertions(+), 42 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 15422da1b3b2e00..3a7885a50412a19 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -14,6 +14,7 @@ import functools import inspect import itertools +import time import types import warnings import weakref @@ -522,60 +523,58 @@ async def _cancel_and_wait(fut, loop): fut.remove_done_callback(cb) -# This is *not* a @coroutine! It is just an iterator (yielding Futures). -def as_completed(fs, *, loop=None, timeout=None): - """Return an iterator whose values are coroutines. +class as_completed(object): + """Asynchronous iterator over awaitables which returns the result of each + as completed. - When waiting for the yielded coroutines you'll get the results (or - exceptions!) of the original Futures (or coroutines), in the order - in which and as soon as they complete. + This differs from PEP 3148; results are yielded from asynchronous + iteration rather than futures. - This differs from PEP 3148; the proper way to use this is: + For backwards compatibility, this can also be used as a regular iterator:: for f in as_completed(fs): result = await f # The 'await' may raise. # Use result. - If a timeout is specified, the 'await' will raise - TimeoutError when the timeout occurs before all Futures are done. - - Note: The futures 'f' are not necessarily members of fs. """ - if futures.isfuture(fs) or coroutines.iscoroutine(fs): - raise TypeError(f"expect a list of futures, not {type(fs).__name__}") - loop = loop if loop is not None else events.get_event_loop() - todo = {ensure_future(f, loop=loop) for f in set(fs)} - from .queues import Queue # Import here to avoid circular import problem. - done = Queue(loop=loop) - timeout_handle = None + def __init__(self, fs, *, loop=None, timeout=None): + from .queue import Queue - def _on_timeout(): - for f in todo: - f.remove_done_callback(_on_completion) - done.put_nowait(None) # Queue a dummy value for _wait_for_one(). - todo.clear() # Can't do todo.remove(f) in the loop. + self._pending = set(ensure_future(f) for f in fs) + self._completed = Queue() - def _on_completion(f): - if not todo: - return # _on_timeout() was here first. - todo.remove(f) - done.put_nowait(f) - if not todo and timeout_handle is not None: - timeout_handle.cancel() + self._loop = loop + self._timeout = timeout - async def _wait_for_one(): - f = await done.get() - if f is None: - # Dummy value from _on_timeout(). - raise exceptions.TimeoutError - return f.result() # May raise f.exception(). + for future in self._pending: + future.add_done_callback(self._done_callback) - for f in todo: - f.add_done_callback(_on_completion) - if todo and timeout is not None: - timeout_handle = loop.call_later(timeout, _on_timeout) - for _ in range(len(todo)): - yield _wait_for_one() + def _done_callback(self, future): + self._pending.remove(future) + self._completed.put_nowait(future) + + async def __aiter__(self): + try: + t0 = time.time() + + while self._pending: + timeout = (self._timeout - (time.time() - t0) + if self._timeout is not None else None) + + future = await wait_for(self._completed.get(), + timeout, + loop=self._loop) + yield future.result() + finally: + # If an exception happened, we want to ensure that the done + # callback doesn't run anyway + for future in self._pending: + future.remove_done_callback(self._done_callback) + + def __iter__(self): + while self._pending: + future = self._completed.get() + yield future @types.coroutine From 67a46d60caf2f13e94e2ec6c0af720e55ac38f10 Mon Sep 17 00:00:00 2001 From: "Michael V. DePalatis" Date: Tue, 30 Oct 2018 20:33:37 -0600 Subject: [PATCH 2/4] Use timeout handling from previous implementation --- Lib/asyncio/tasks.py | 49 ++++++++++++++++++++++++++++++-------------- 1 file changed, 34 insertions(+), 15 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 3a7885a50412a19..a4819dd0c77f86c 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -14,7 +14,6 @@ import functools import inspect import itertools -import time import types import warnings import weakref @@ -545,36 +544,56 @@ def __init__(self, fs, *, loop=None, timeout=None): self._loop = loop self._timeout = timeout + self._timeout_handle = None for future in self._pending: - future.add_done_callback(self._done_callback) + future.add_done_callback(self._on_completion) - def _done_callback(self, future): + def _on_completion(self, future): self._pending.remove(future) self._completed.put_nowait(future) + if not self._pending and self._timeout_handle is not None: + self._timeout_handle.cancel() + + async def _wait_for_one(self): + f = await self._completed.get() + if f is None: + # Dummy value from _on_timeout(). + raise exceptions.TimeoutError + + return f.result() # May raise f.exception(). + + def _on_timeout(self): + for f in self._pending: + f.remove_done_callback(self._on_completion) + + def _start_timeout(self): + if self._pending and self._timeout is not None: + self._timeout_handle = self._loop.call_later(self._timeout, + self._on_timeout) + async def __aiter__(self): - try: - t0 = time.time() + self._start_timeout() + try: while self._pending: - timeout = (self._timeout - (time.time() - t0) - if self._timeout is not None else None) - - future = await wait_for(self._completed.get(), - timeout, - loop=self._loop) + future = await self._completed.get() yield future.result() finally: # If an exception happened, we want to ensure that the done # callback doesn't run anyway for future in self._pending: - future.remove_done_callback(self._done_callback) + future.remove_done_callback(self._on_completion) def __iter__(self): - while self._pending: - future = self._completed.get() - yield future + for f in self._pending: + f.add_done_callback(self._on_completion) + + self._start_timeout() + + for _ in range(len(self._pending)): + yield self._wait_for_one() @types.coroutine From 52b10dade6361baa6cb0bb030e2c060591d87d4f Mon Sep 17 00:00:00 2001 From: "Michael V. DePalatis" Date: Wed, 31 Oct 2018 08:46:26 -0600 Subject: [PATCH 3/4] Fix import --- Lib/asyncio/tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index a4819dd0c77f86c..ad23fd7c370a56f 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -537,7 +537,7 @@ class as_completed(object): """ def __init__(self, fs, *, loop=None, timeout=None): - from .queue import Queue + from .queues import Queue self._pending = set(ensure_future(f) for f in fs) self._completed = Queue() From ab8e7627451abf5040390c9c922bf62eeef09292 Mon Sep 17 00:00:00 2001 From: "Michael V. DePalatis" Date: Wed, 31 Oct 2018 18:37:28 -0600 Subject: [PATCH 4/4] Yield a future with async for --- Lib/asyncio/tasks.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index ad23fd7c370a56f..8e3149e7c16d799 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -523,11 +523,11 @@ async def _cancel_and_wait(fut, loop): class as_completed(object): - """Asynchronous iterator over awaitables which returns the result of each - as completed. + """Asynchronous iterator over awaitables which yields each future as it + is completed consistent with PEP 3148. Usage:: - This differs from PEP 3148; results are yielded from asynchronous - iteration rather than futures. + async for future in as_completed(fs): + print(future.result()) For backwards compatibility, this can also be used as a regular iterator:: @@ -579,7 +579,7 @@ async def __aiter__(self): try: while self._pending: future = await self._completed.get() - yield future.result() + yield future finally: # If an exception happened, we want to ensure that the done # callback doesn't run anyway