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 ad31f5d..f486b67 100644
--- 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 74fc1e4..01f62b7 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -1548,6 +1548,30 @@ def gen():
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():