python · Dreamsorcerer · Sep 3, 2021 · Sep 3, 2021 · Sep 3, 2021 · Sep 3, 2021
@@ -25,6 +25,8 @@
 from . import futures
 from .coroutines import _is_coroutine
 
+_SENTINEL = object()
+
 # Helper to generate new task names
 # This uses itertools.count() instead of a "+= 1" operation because the latter
 # is not thread safe. See bpo-11866 for a longer explanation.
@@ -432,8 +434,8 @@ async def wait_for(fut, timeout):
         # wait until the future completes or the timeout
         try:
             await waiter
-        except exceptions.CancelledError:
-            if fut.done():
+        except exceptions.CancelledError as e:
+            if fut.done() and _SENTINEL in e.args:
                 return fut.result()
             else:
                 fut.remove_done_callback(cb)
@@ -455,11 +457,9 @@ async def wait_for(fut, timeout):
             # exception, we should re-raise it
             # See https://bugs.python.org/issue40607
             try:
-                fut.result()
+                return fut.result()
             except exceptions.CancelledError as exc:
                 raise exceptions.TimeoutError() from exc
-            else:
-                raise exceptions.TimeoutError()
     finally:
         timeout_handle.cancel()
 
@@ -516,7 +516,7 @@ async def _cancel_and_wait(fut, loop):
     fut.add_done_callback(cb)
 
     try:
-        fut.cancel()
+        fut.cancel(_SENTINEL)
         # We cannot wait on *fut* directly to make
         # sure _cancel_and_wait itself is reliably cancellable.
         await waiter

@@ -1009,20 +1009,33 @@ 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
+        async def inner():
+            await asyncio.wait_for(asyncio.sleep(1), timeout=2)
+            return 1
 
-        loop = self.new_test_loop(gen)
+        async def main():
+            result = await asyncio.wait_for(inner(), timeout=1)
+            assert result == 1
 
-        fut = self.new_future(loop)
-        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")
+        asyncio.run(main())
+
+    def test_wait_for_does_not_suppress_cancellation(self):
+        async def inner():
+            return
+
+        async def with_for_coro():
+            await asyncio.wait_for(inner(), timeout=1)
+            assert False, 'End of with_for_coro. Should not be reached!'
+
+        async def main():
+            task = asyncio.create_task(with_for_coro())
+            await asyncio.sleep(0)
+            self.assertFalse(task.done())
+            task.cancel()
+            with self.assertRaises(asyncio.CancelledError):
+                await task
+
+        asyncio.run(main())
 
     def test_wait_for_waits_for_task_cancellation(self):
         loop = asyncio.new_event_loop()
@@ -1101,24 +1114,6 @@ async def inner():
         with self.assertRaises(FooException):
             loop.run_until_complete(foo())
 
-    def test_wait_for_raises_timeout_error_if_returned_during_cancellation(self):
-        loop = asyncio.new_event_loop()
-        self.addCleanup(loop.close)
-
-        async def foo():
-            async def inner():
-                try:
-                    await asyncio.sleep(0.2)
-                except asyncio.CancelledError:
-                    return 42
-
-            inner_task = self.new_task(loop, inner())
-
-            await asyncio.wait_for(inner_task, timeout=_EPSILON)
-
-        with self.assertRaises(asyncio.TimeoutError):
-            loop.run_until_complete(foo())
-
     def test_wait_for_self_cancellation(self):
         loop = asyncio.new_event_loop()
         self.addCleanup(loop.close)