From 620279b9ace3fff66245672bf7efbb62b2969a30 Mon Sep 17 00:00:00 2001 From: Yury Selivanov Date: Fri, 2 Oct 2015 15:00:19 -0400 Subject: asyncio: ensure_future() now understands awaitables --- Lib/asyncio/tasks.py | 16 ++++++++++++++-- Lib/test/test_asyncio/test_tasks.py | 18 ++++++++++++++++++ 2 files changed, 32 insertions(+), 2 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index a235e74..434f498 100644 --- 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 index 0426787..16d3d9d 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -153,6 +153,24 @@ class TaskTests(test_utils.TestCase): 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') -- cgit v0.12