bpo-42130: Fix for explicit suppressing of cancellations in wait_for()

Lib/asyncio/tasks.py

@@ -439,7 +439,7 @@ async def wait_for(fut, timeout):
  # after wait_for() returns.
  # See https://bugs.python.org/issue32751
  await _cancel_and_wait(fut, loop=loop)
-raise
+return fut.result()
 
  if fut.done():
  return fut.result()

Lib/test/test_asyncio/test_tasks.py

@@ -1009,16 +1009,63 @@ def gen():
  self.assertEqual(res, "ok")
 
  def test_wait_for_cancellation_race_condition(self):
+ def gen():
+ yield 0.1
+ yield 0.1
+ yield 0.1
+ yield 0.1
+
+ loop = self.new_test_loop(gen)
+
+ fut = self.new_future(loop)
+ # Test that if task is cancelled at the same time the future
+ # completes, the result is still returned.
+ loop.call_later(0.1, fut.set_result, "ok")
+ task = loop.create_task(asyncio.wait_for(fut, timeout=1))
+ loop.call_later(0.1, task.cancel)
+ res = loop.run_until_complete(task)
+ self.assertEqual(res, "ok")
+
+ def test_wait_for_suppress_cancellation(self):
+ def gen():
+ yield
+ yield 0.1
+ yield 0.1
+ yield 0.1
+
  async def inner():
-with contextlib.suppress(asyncio.CancelledError):
+try:
  await asyncio.sleep(1)
- return 1
+ except asyncio.CancelledError:
+ return "ok"
 
- async def main():
- result = await asyncio.wait_for(inner(), timeout=.01)
- assert result == 1
+ loop = self.new_test_loop(gen)
+
+ fut = self.new_future(loop)
+ task = loop.create_task(asyncio.wait_for(inner(), timeout=1))
+ loop.call_later(0.1, task.cancel)
+ # Cancellation is suppressed in inner(), so should still return here.
+ res = loop.run_until_complete(task)
+ self.assertEqual(res, "ok")
+
+ def test_wait_for_cancellation_inner_race_condition(self):
+ def gen():
+ yield
+ yield 0.1
+ yield 0
+ yield 0
+
+ async def inner():
+ # Test that if fut completes at same time as timeout, the result
+ # is still returned.
+ return await asyncio.wait_for(fut, timeout=1)
 
- asyncio.run(main())
+ loop = self.new_test_loop(gen)
+
+ fut = self.new_future(loop)
+ loop.call_later(0.1, fut.set_result, "ok")
+ res = loop.run_until_complete(asyncio.wait_for(inner(), timeout=0.1))
+ self.assertEqual(res, "ok")
 
  def test_wait_for_waits_for_task_cancellation(self):
  loop = asyncio.new_event_loop()

Misc/NEWS.d/next/Library/2021-11-29-17-17-45.bpo-42130.Ho8dUh.rst

@@ -0,0 +1,2 @@
+Fix issue where ``asyncio.wait_for()`` is cancelled when the coro
+actually suppresses the cancellation.