GH-78530: add support for generators in `asyncio.wait` (#102761)

4 files changed, 24 insertions, 0 deletions

diff --git a/Doc/library/asyncio-task.rst b/Doc/library/asyncio-task.rst
index 5f1449e..a0900cd 100644
--- a/Doc/library/asyncio-task.rst
+++ b/Doc/library/asyncio-task.rst
@@ -804,6 +804,10 @@ Waiting Primitives
    .. versionchanged:: 3.11
       Passing coroutine objects to ``wait()`` directly is forbidden.
 
+   .. versionchanged:: 3.12
+      Added support for generators yielding tasks.
+
+
 .. function:: as_completed(aws, *, timeout=None)
 
    Run :ref:`awaitable objects <asyncio-awaitables>` in the *aws*
diff --git a/Doc/whatsnew/3.12.rst b/Doc/whatsnew/3.12.rst
index a398cd9..b55b961 100644
--- a/Doc/whatsnew/3.12.rst
+++ b/Doc/whatsnew/3.12.rst
@@ -241,6 +241,9 @@ asyncio
   :mod:`asyncio` does not support legacy generator-based coroutines.
   (Contributed by Kumar Aditya in :gh:`102748`.)
 
+* :func:`asyncio.wait` now accepts generators yielding tasks.
+  (Contributed by Kumar Aditya in :gh:`78530`.)
+
 inspect
 -------
 
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 5b935b5..731fa0c 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1373,6 +1373,22 @@ class BaseTaskTests:
         self.assertEqual(res, 42)
         self.assertAlmostEqual(0.15, loop.time())
 
+
+    def test_wait_generator(self):
+        async def func(a):
+            return a
+
+        loop = self.new_test_loop()
+
+        async def main():
+            tasks = (self.new_task(loop, func(i)) for i in range(10))
+            done, pending = await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED)
+            self.assertEqual(len(done), 10)
+            self.assertEqual(len(pending), 0)
+
+        loop.run_until_complete(main())
+
+
     def test_as_completed(self):
 
         def gen():
diff --git a/Misc/NEWS.d/next/Library/2023-03-16-16-43-04.gh-issue-78530.Lr8eq_.rst b/Misc/NEWS.d/next/Library/2023-03-16-16-43-04.gh-issue-78530.Lr8eq_.rst
new file mode 100644
index 0000000..bdb46d0
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2023-03-16-16-43-04.gh-issue-78530.Lr8eq_.rst
@@ -0,0 +1 @@
+:func:`asyncio.wait` now accepts generators yielding tasks. Patch by Kumar Aditya.