diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index 42413dc00f0..b52933fc6c6 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -382,8 +382,9 @@ def wait_for(fut, timeout, *, loop=None): Coroutine will be wrapped in Task. - Returns result of the Future or coroutine. Raises TimeoutError when - timeout occurs. + Returns result of the Future or coroutine. When a timeout occurs, + it cancels the task and raises TimeoutError. To avoid the task + cancellation, wrap it in shield(). Usage: @@ -405,6 +406,7 @@ def wait_for(fut, timeout, *, loop=None): return fut.result() else: fut.remove_done_callback(cb) + fut.cancel() raise futures.TimeoutError() finally: timeout_handle.cancel() diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py index 79a25d298fc..3d08ad8be8b 100644 --- a/Lib/test/test_asyncio/test_tasks.py +++ b/Lib/test/test_asyncio/test_tasks.py @@ -355,30 +355,32 @@ def gen(): when = yield 0 self.assertAlmostEqual(0.1, when) when = yield 0.1 - self.assertAlmostEqual(0.4, when) - yield 0.1 loop = test_utils.TestLoop(gen) self.addCleanup(loop.close) + foo_running = None + @tasks.coroutine def foo(): - yield from tasks.sleep(0.2, loop=loop) + nonlocal foo_running + foo_running = True + try: + yield from tasks.sleep(0.2, loop=loop) + finally: + foo_running = False return 'done' fut = tasks.Task(foo(), loop=loop) with self.assertRaises(futures.TimeoutError): loop.run_until_complete(tasks.wait_for(fut, 0.1, loop=loop)) - - self.assertFalse(fut.done()) + self.assertTrue(fut.done()) + # it should have been cancelled due to the timeout + self.assertTrue(fut.cancelled()) self.assertAlmostEqual(0.1, loop.time()) + self.assertEqual(foo_running, False) - # wait for result - res = loop.run_until_complete( - tasks.wait_for(fut, 0.3, loop=loop)) - self.assertEqual(res, 'done') - self.assertAlmostEqual(0.2, loop.time()) def test_wait_for_with_global_loop(self): @@ -406,11 +408,8 @@ def foo(): events.set_event_loop(None) self.assertAlmostEqual(0.01, loop.time()) - self.assertFalse(fut.done()) - - # move forward to close generator - loop.advance_time(10) - loop.run_until_complete(fut) + self.assertTrue(fut.done()) + self.assertTrue(fut.cancelled()) def test_wait(self):