asyncio: Make ensure_future() accept all kinds of awaitables.

1st1 · 1st1 · commit 43d71e251270 · 2015-10-02T15:05:59.000-04:00
diff --git a/Doc/whatsnew/3.5.rst b/Doc/whatsnew/3.5.rst
@@ -803,6 +803,13 @@ Notable changes in the :mod:`asyncio` module since Python 3.4.0:
   :class:`asyncio.Queue` class.
   (Contributed by Victor Stinner.)
 
+Updates in 3.5.1:
+
+* The :func:`~asyncio.ensure_future` function and all functions that
+  use it, such as :meth:`loop.run_until_complete() <asyncio.BaseEventLoop.run_until_complete>`,
+  now accept all kinds of :term:`awaitable objects <awaitable>`.
+  (Contributed by Yury Selivanov.)
+
 
 bz2
 ---
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -512,7 +512,7 @@ def async(coro_or_future, *, loop=None):
 
 
 def ensure_future(coro_or_future, *, loop=None):
-    """Wrap a coroutine in a future.
+    """Wrap a coroutine or an awaitable in a future.
 
     If the argument is a Future, it is returned directly.
     """
@@ -527,8 +527,20 @@ def ensure_future(coro_or_future, *, loop=None):
         if task._source_traceback:
             del task._source_traceback[-1]
         return task
+    elif compat.PY35 and inspect.isawaitable(coro_or_future):
+        return ensure_future(_wrap_awaitable(coro_or_future), loop=loop)
     else:
-        raise TypeError('A Future or coroutine is required')
+        raise TypeError('A Future, a coroutine or an awaitable is required')
+
+
+@coroutine
+def _wrap_awaitable(awaitable):
+    """Helper for asyncio.ensure_future().
+
+    Wraps awaitable (an object with __await__) into a coroutine
+    that will later be wrapped in a Task by ensure_future().
+    """
+    return (yield from awaitable.__await__())
 
 
 class _GatheringFuture(futures.Future):
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
@@ -153,6 +153,24 @@ def notmuch():
         t = asyncio.ensure_future(t_orig, loop=self.loop)
         self.assertIs(t, t_orig)
 
+    @unittest.skipUnless(PY35, 'need python 3.5 or later')
+    def test_ensure_future_awaitable(self):
+        class Aw:
+            def __init__(self, coro):
+                self.coro = coro
+            def __await__(self):
+                return (yield from self.coro)
+
+        @asyncio.coroutine
+        def coro():
+            return 'ok'
+
+        loop = asyncio.new_event_loop()
+        self.set_event_loop(loop)
+        fut = asyncio.ensure_future(Aw(coro()), loop=loop)
+        loop.run_until_complete(fut)
+        assert fut.result() == 'ok'
+
     def test_ensure_future_neither(self):
         with self.assertRaises(TypeError):
             asyncio.ensure_future('ok')
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -152,6 +152,8 @@ Library
 - Issue #23572: Fixed functools.singledispatch on classes with falsy
   metaclasses.  Patch by Ethan Furman.
 
+- asyncio: ensure_future() now accepts awaitable objects.
+
 IDLE
 ----
 
