diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 7959a55afea..63412a97e62 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/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/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 7807dc04719..06447d778b5 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/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: