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

Author: kumaraditya303

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*

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
 -------
 

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():

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.