From 560c6beda55e96ba3dde8f5bdc9f7685502bcf0c Mon Sep 17 00:00:00 2001
From: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>
Date: Fri, 17 Mar 2023 06:58:43 +0530
Subject: [PATCH] GH-78530: add support for generators in `asyncio.wait`
 (#102761)

---
 Doc/library/asyncio-task.rst                     |  4 ++++
 Doc/whatsnew/3.12.rst                            |  3 +++
 Lib/test/test_asyncio/test_tasks.py              | 16 ++++++++++++++++
 ...2023-03-16-16-43-04.gh-issue-78530.Lr8eq_.rst |  1 +
 4 files changed, 24 insertions(+)
 create mode 100644 Misc/NEWS.d/next/Library/2023-03-16-16-43-04.gh-issue-78530.Lr8eq_.rst

diff --git a/Doc/library/asyncio-task.rst b/Doc/library/asyncio-task.rst
index 5f1449e1b599efb..a0900cd25a7731b 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 9e095569e15562c..242d76e969ba977 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 5b935b526541a1c..731fa0c5a60b9b9 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1373,6 +1373,22 @@ async def foo():
         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 000000000000000..bdb46d08c5c4af4
--- /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.