diff options
Diffstat (limited to 'Lib/asyncio/tasks.py')
-rw-r--r-- | Lib/asyncio/tasks.py | 16 |
1 files changed, 14 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): |