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

This commit is contained in:
Kumar Aditya 2023-03-16 09:20:43 +05:30 committed by GitHub
parent 1c9f3391b9
commit a44553ea9f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 22 additions and 16 deletions

View File

@ -25,7 +25,6 @@
from . import exceptions from . import exceptions
from . import futures from . import futures
from . import timeouts from . import timeouts
from .coroutines import _is_coroutine
# Helper to generate new task names # Helper to generate new task names
# This uses itertools.count() instead of a "+= 1" operation because the latter # This uses itertools.count() instead of a "+= 1" operation because the latter
@ -635,11 +634,14 @@ def ensure_future(coro_or_future, *, loop=None):
raise ValueError('The future belongs to a different loop than ' raise ValueError('The future belongs to a different loop than '
'the one specified as the loop argument') 'the one specified as the loop argument')
return coro_or_future return coro_or_future
called_wrap_awaitable = False should_close = True
if not coroutines.iscoroutine(coro_or_future): if not coroutines.iscoroutine(coro_or_future):
if inspect.isawaitable(coro_or_future): if inspect.isawaitable(coro_or_future):
async def _wrap_awaitable(awaitable):
return await awaitable
coro_or_future = _wrap_awaitable(coro_or_future) coro_or_future = _wrap_awaitable(coro_or_future)
called_wrap_awaitable = True should_close = False
else: else:
raise TypeError('An asyncio.Future, a coroutine or an awaitable ' raise TypeError('An asyncio.Future, a coroutine or an awaitable '
'is required') 'is required')
@ -649,23 +651,11 @@ def ensure_future(coro_or_future, *, loop=None):
try: try:
return loop.create_task(coro_or_future) return loop.create_task(coro_or_future)
except RuntimeError: except RuntimeError:
if not called_wrap_awaitable: if should_close:
coro_or_future.close() coro_or_future.close()
raise 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): class _GatheringFuture(futures.Future):
"""Helper for gather(). """Helper for gather().

View File

@ -8,6 +8,7 @@
import re import re
import sys import sys
import traceback import traceback
import types
import unittest import unittest
from unittest import mock from unittest import mock
from types import GenericAlias from types import GenericAlias
@ -274,6 +275,20 @@ async def coro():
loop.run_until_complete(fut) loop.run_until_complete(fut)
self.assertEqual(fut.result(), 'ok') 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): def test_ensure_future_neither(self):
with self.assertRaises(TypeError): with self.assertRaises(TypeError):
asyncio.ensure_future('ok') asyncio.ensure_future('ok')

View File

@ -0,0 +1 @@
:meth:`asyncio.Task.get_coro` now always returns a coroutine when wrapping an awaitable object. Patch by Kumar Aditya.