python · mivade · Oct 31, 2018 · Oct 31, 2018 · Oct 31, 2018 · Nov 1, 2018
@@ -570,62 +570,79 @@ async def _cancel_and_wait(fut):
         fut.remove_done_callback(cb)
 
 
-# This is *not* a @coroutine!  It is just an iterator (yielding Futures).
-def as_completed(fs, *, timeout=None):
-    """Return an iterator whose values are coroutines.
+class as_completed(object):
-class as_completed(object):
+class as_completed:
-class as_completed(object):
+class as_completed:
+    """Asynchronous iterator over awaitables which yields each future as it
+    is completed consistent with PEP 3148. Usage::
 
-    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.
+        async for future in as_completed(fs):
+            print(future.result())
 
-    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 an iterable of futures, not {type(fs).__name__}")
+    def __init__(self, fs, *, timeout=None):
+        if futures.isfuture(fs) or coroutines.iscoroutine(fs):
+            raise TypeError(f"expect an iterable of futures, not {type(fs).__name__}")
 
-    from .queues import Queue  # Import here to avoid circular import problem.
-    done = Queue()
+        from .queues import Queue  # Import here to avoid circular import problem.
+        self._loop = events.get_event_loop()
+        self._pending = set(ensure_future(f, loop=self._loop) for f in fs)
+        self._completed = Queue()
+        self._timeout = timeout
+        self._timeout_handle = None
 
-    loop = events.get_event_loop()
-    todo = {ensure_future(f, loop=loop) for f in set(fs)}
-    timeout_handle = None
+        for future in self._pending:
+            future.add_done_callback(self._on_completion)
 
-    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.
+    def _on_completion(self, future):
+        self._pending.remove(future)
+        self._completed.put_nowait(future)
 
-    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()
+        if not self._pending and self._timeout_handle is not None:
+            self._timeout_handle.cancel()
 
-    async def _wait_for_one():
-        f = await done.get()
+    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().
 
-    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 _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):
+        self._start_timeout()
+
+        try:
+            while self._pending:
+                future = await self._completed.get()
+                yield future
+        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._on_completion)
+
+    def __iter__(self):
+        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