bpo-42140: Improve asyncio.wait function (GH-22938)

# Improve asyncio.wait function

The original code creates the futures set two times.
We can create this set before, avoiding the second creation.

This new behaviour [breaks the aiokafka library](https://github.com/aio-libs/aiokafka/pull/672), because it gives an iterator to that function, so the second iteration become empty.

Automerge-Triggered-By: GH:1st1

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index ad31f5d59796b11ec647e3ac01e404ad14f1f9e9..f486b6722941119b896d97412ec425dbe2a7a9db 100644 (file)
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -400,13 +400,15 @@ async def wait(fs, *, loop=None, timeout=None, return_when=ALL_COMPLETED):
                       "and scheduled for removal in Python 3.10.",
                       DeprecationWarning, stacklevel=2)
 
-    if any(coroutines.iscoroutine(f) for f in set(fs)):
+    fs = set(fs)
+
+    if any(coroutines.iscoroutine(f) for f in fs):
         warnings.warn("The explicit passing of coroutine objects to "
                       "asyncio.wait() is deprecated since Python 3.8, and "
                       "scheduled for removal in Python 3.11.",
                       DeprecationWarning, stacklevel=2)
 
-    fs = {ensure_future(f, loop=loop) for f in set(fs)}
+    fs = {ensure_future(f, loop=loop) for f in fs}
 
     return await _wait(fs, timeout, return_when, loop)
 

diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index 74fc1e4a42133c30c5123d969bbcff2da2cd293a..01f62b7f408726435ce538b5cd4ecab8e741a696 100644 (file)
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1548,6 +1548,30 @@ class BaseTaskTests:
         loop.advance_time(10)
         loop.run_until_complete(asyncio.wait([a, b]))
 
+    def test_wait_with_iterator_of_tasks(self):
+
+        def gen():
+            when = yield
+            self.assertAlmostEqual(0.1, when)
+            when = yield 0
+            self.assertAlmostEqual(0.15, when)
+            yield 0.15
+
+        loop = self.new_test_loop(gen)
+
+        a = self.new_task(loop, asyncio.sleep(0.1))
+        b = self.new_task(loop, asyncio.sleep(0.15))
+
+        async def foo():
+            done, pending = await asyncio.wait(iter([b, a]))
+            self.assertEqual(done, set([a, b]))
+            self.assertEqual(pending, set())
+            return 42
+
+        res = loop.run_until_complete(self.new_task(loop, foo()))
+        self.assertEqual(res, 42)
+        self.assertAlmostEqual(0.15, loop.time())
+
     def test_as_completed(self):
 
         def gen():

diff --git a/Misc/NEWS.d/next/Library/2020-10-24-04-02-36.bpo-42140.miLqvb.rst b/Misc/NEWS.d/next/Library/2020-10-24-04-02-36.bpo-42140.miLqvb.rst
new file mode 100644 (file)
index 0000000..4160234
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-10-24-04-02-36.bpo-42140.miLqvb.rst
@@ -0,0 +1 @@
+Improve asyncio.wait function to create the futures set just one time.