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.