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

2 files changed, 21 insertions, 16 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 1c20754..c90d32c 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -25,7 +25,6 @@ from . import events
 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
@@ -635,11 +634,14 @@ 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 = True
     if not coroutines.iscoroutine(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)
-            called_wrap_awaitable = True
+            should_close = False
         else:
             raise TypeError('An asyncio.Future, a coroutine or an awaitable '
                             'is required')
@@ -649,23 +651,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().
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index e533d52..5b935b5 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -8,6 +8,7 @@ import random
 import re
 import sys
 import traceback
+import types
 import unittest
 from unittest import mock
 from types import GenericAlias
@@ -274,6 +275,20 @@ class BaseTaskTests:
         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')