Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GH-100112: avoid using iterable coroutines in asyncio internally #100128

Merged
merged 5 commits into from
Mar 16, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 9 additions & 16 deletions Lib/asyncio/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
from . import exceptions
from . import futures
from . import timeouts
from .coroutines import _is_coroutine

# Helper to generate new task names
# This uses itertools.count() instead of a "+= 1" operation because the latter
Expand Down Expand Up @@ -635,11 +634,17 @@ def ensure_future(coro_or_future, *, loop=None):
raise ValueError('The future belongs to a different loop than '
'the one specified as the loop argument')
return coro_or_future
called_wrap_awaitable = False
should_close = False
if not coroutines.iscoroutine(coro_or_future):
if inspect.isawaitable(coro_or_future):
async def _wrap_awaitable(awaitable):
@types.coroutine
def wrapper():
return (yield from awaitable.__await__())
return await wrapper()
kumaraditya303 marked this conversation as resolved.
Show resolved Hide resolved

coro_or_future = _wrap_awaitable(coro_or_future)
called_wrap_awaitable = True
should_close = True
else:
raise TypeError('An asyncio.Future, a coroutine or an awaitable '
'is required')
Expand All @@ -649,23 +654,11 @@ def ensure_future(coro_or_future, *, loop=None):
try:
return loop.create_task(coro_or_future)
except RuntimeError:
if not called_wrap_awaitable:
if should_close:
coro_or_future.close()
raise


@types.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__())

_wrap_awaitable._is_coroutine = _is_coroutine


class _GatheringFuture(futures.Future):
"""Helper for gather().

Expand Down
15 changes: 15 additions & 0 deletions Lib/test/test_asyncio/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import re
import sys
import traceback
import types
import unittest
from unittest import mock
from types import GenericAlias
Expand Down Expand Up @@ -274,6 +275,20 @@ async def coro():
loop.run_until_complete(fut)
self.assertEqual(fut.result(), 'ok')

def test_ensure_future_task_awaitable(self):
class Aw:
def __await__(self):
return asyncio.sleep(0, result='ok').__await__()

loop = asyncio.new_event_loop()
self.set_event_loop(loop)
task = asyncio.ensure_future(Aw(), loop=loop)
loop.run_until_complete(task)
self.assertTrue(task.done())
self.assertEqual(task.result(), 'ok')
self.assertIsInstance(task.get_coro(), types.CoroutineType)
loop.close()

def test_ensure_future_neither(self):
with self.assertRaises(TypeError):
asyncio.ensure_future('ok')
Expand Down