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

GH- 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](aio-libs/aiokafka#672), because it gives an iterator to that function, so the second iteration become empty.

Automerge-Triggered-By: GH:1st1
(cherry picked from commit 7e5ef0a)

Co-authored-by: Diogo Dutra <diogodutradamata@gmail.com>

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)
 

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

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.