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

2 files changed, 35 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()
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
index 7807dc0..06447d7 100644
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -1705,6 +1705,33 @@ class TaskTests(test_utils.TestCase):
                           'test_task_source_traceback'))
         self.loop.run_until_complete(task)
 
+    def _test_cancel_wait_for(self, timeout):
+        loop = asyncio.new_event_loop()
+        self.addCleanup(loop.close)
+
+        @asyncio.coroutine
+        def blocking_coroutine():
+            fut = asyncio.Future(loop=loop)
+            # Block: fut result is never set
+            yield from fut
+
+        task = loop.create_task(blocking_coroutine())
+
+        wait = loop.create_task(asyncio.wait_for(task, timeout, loop=loop))
+        loop.call_soon(wait.cancel)
+
+        self.assertRaises(asyncio.CancelledError,
+                          loop.run_until_complete, wait)
+
+        # Python issue #23219: cancelling the wait must also cancel the task
+        self.assertTrue(task.cancelled())
+
+    def test_cancel_blocking_wait_for(self):
+        self._test_cancel_wait_for(None)
+
+    def test_cancel_wait_for(self):
+        self._test_cancel_wait_for(60.0)
+
 
 class GatherTestsBase: