Python issue #23219: cancelling wait_for() now cancels the task

1 files changed, 8 insertions, 4 deletions

diff --git a/asyncio/tasks.py b/asyncio/tasks.py
index 7959a55..63412a9 100644
--- a/asyncio/tasks.py
+++ b/asyncio/tasks.py
@@ -347,10 +347,9 @@ def wait_for(fut, timeout, *, loop=None):
     it cancels the task and raises TimeoutError.  To avoid the task
     cancellation, wrap it in shield().
 
-    Usage:
-
-        result = yield from asyncio.wait_for(fut, 10.0)
+    If the wait is cancelled, the task is also cancelled.
 
+    This function is a coroutine.
     """
     if loop is None:
         loop = events.get_event_loop()
@@ -367,7 +366,12 @@ def wait_for(fut, timeout, *, loop=None):
 
     try:
         # wait until the future completes or the timeout
-        yield from waiter
+        try:
+            yield from waiter
+        except futures.CancelledError:
+            fut.remove_done_callback(cb)
+            fut.cancel()
+            raise
 
         if fut.done():
             return fut.result()