diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index e78719de216fd0..a2e06d5ef72f42 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -24,6 +24,7 @@ from . import events from . import exceptions from . import futures +from . import timeouts from .coroutines import _is_coroutine # Helper to generate new task names @@ -437,65 +438,44 @@ async def wait_for(fut, timeout): If the wait is cancelled, the task is also cancelled. + If the task supresses the cancellation and returns a value instead, + that value is returned. + This function is a coroutine. """ - loop = events.get_running_loop() + # The special case for timeout <= 0 is for the following case: + # + # async def test_waitfor(): + # func_started = False + # + # async def func(): + # nonlocal func_started + # func_started = True + # + # try: + # await asyncio.wait_for(func(), 0) + # except asyncio.TimeoutError: + # assert not func_started + # else: + # assert False + # + # asyncio.run(test_waitfor()) - if timeout is None: - return await fut - if timeout <= 0: - fut = ensure_future(fut, loop=loop) + if timeout is not None and timeout <= 0: + fut = ensure_future(fut) if fut.done(): return fut.result() - await _cancel_and_wait(fut, loop=loop) + await _cancel_and_wait(fut) try: return fut.result() except exceptions.CancelledError as exc: - raise exceptions.TimeoutError() from exc - - waiter = loop.create_future() - timeout_handle = loop.call_later(timeout, _release_waiter, waiter) - cb = functools.partial(_release_waiter, waiter) - - fut = ensure_future(fut, loop=loop) - fut.add_done_callback(cb) - - try: - # wait until the future completes or the timeout - try: - await waiter - except exceptions.CancelledError: - if fut.done(): - return fut.result() - else: - fut.remove_done_callback(cb) - # We must ensure that the task is not running - # after wait_for() returns. - # See https://bugs.python.org/issue32751 - await _cancel_and_wait(fut, loop=loop) - raise - - if fut.done(): - return fut.result() - else: - fut.remove_done_callback(cb) - # We must ensure that the task is not running - # after wait_for() returns. - # See https://bugs.python.org/issue32751 - await _cancel_and_wait(fut, loop=loop) - # In case task cancellation failed with some - # exception, we should re-raise it - # See https://bugs.python.org/issue40607 - try: - return fut.result() - except exceptions.CancelledError as exc: - raise exceptions.TimeoutError() from exc - finally: - timeout_handle.cancel() + raise TimeoutError from exc + async with timeouts.timeout(timeout): + return await fut async def _wait(fs, timeout, return_when, loop): """Internal helper for wait(). @@ -541,9 +521,10 @@ def _on_completion(f): return done, pending -async def _cancel_and_wait(fut, loop): +async def _cancel_and_wait(fut): """Cancel the *fut* future or task and wait until it completes.""" + loop = events.get_running_loop() waiter = loop.create_future() cb = functools.partial(_release_waiter, waiter) fut.add_done_callback(cb) diff --git a/Lib/test/test_asyncio/test_futures2.py b/Lib/test/test_asyncio/test_futures2.py index 9e7a5775a70383..b7cfffb76bd8f1 100644 --- a/Lib/test/test_asyncio/test_futures2.py +++ b/Lib/test/test_asyncio/test_futures2.py @@ -86,10 +86,9 @@ async def test_recursive_repr_for_pending_tasks(self): async def func(): return asyncio.all_tasks() - # The repr() call should not raise RecursiveError at first. - # The check for returned string is not very reliable but - # exact comparison for the whole string is even weaker. - self.assertIn('...', repr(await asyncio.wait_for(func(), timeout=10))) + # The repr() call should not raise RecursionError at first. + waiter = await asyncio.wait_for(asyncio.Task(func()),timeout=10) + self.assertIn('...', repr(waiter)) if __name__ == '__main__': diff --git a/Lib/test/test_asyncio/test_waitfor.py b/Lib/test/test_asyncio/test_waitfor.py index 45498fa097f6bc..ed80540b2b3852 100644 --- a/Lib/test/test_asyncio/test_waitfor.py +++ b/Lib/test/test_asyncio/test_waitfor.py @@ -237,33 +237,6 @@ async def inner(): with self.assertRaises(FooException): await foo() - async def test_wait_for_self_cancellation(self): - async def inner(): - try: - await asyncio.sleep(0.3) - except asyncio.CancelledError: - try: - await asyncio.sleep(0.3) - except asyncio.CancelledError: - await asyncio.sleep(0.3) - - return 42 - - inner_task = asyncio.create_task(inner()) - - wait = asyncio.wait_for(inner_task, timeout=0.1) - - # Test that wait_for itself is properly cancellable - # even when the initial task holds up the initial cancellation. - task = asyncio.create_task(wait) - await asyncio.sleep(0.2) - task.cancel() - - with self.assertRaises(asyncio.CancelledError): - await task - - self.assertEqual(await inner_task, 42) - async def _test_cancel_wait_for(self, timeout): loop = asyncio.get_running_loop() @@ -289,6 +262,106 @@ async def test_cancel_blocking_wait_for(self): async def test_cancel_wait_for(self): await self._test_cancel_wait_for(60.0) + async def test_wait_for_cancel_suppressed(self): + # GH-86296: Supressing CancelledError is discouraged + # but if a task subpresses CancelledError and returns a value, + # `wait_for` should return the value instead of raising CancelledError. + # This is the same behavior as `asyncio.timeout`. + + async def return_42(): + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + return 42 + + res = await asyncio.wait_for(return_42(), timeout=0.1) + self.assertEqual(res, 42) + + + async def test_wait_for_issue86296(self): + # GH-86296: The task should get cancelled and not run to completion. + # inner completes in one cycle of the event loop so it + # completes before the task is cancelled. + + async def inner(): + return 'done' + + inner_task = asyncio.create_task(inner()) + reached_end = False + + async def wait_for_coro(): + await asyncio.wait_for(inner_task, timeout=100) + await asyncio.sleep(1) + nonlocal reached_end + reached_end = True + + task = asyncio.create_task(wait_for_coro()) + self.assertFalse(task.done()) + # Run the task + await asyncio.sleep(0) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + self.assertTrue(inner_task.done()) + self.assertEqual(await inner_task, 'done') + self.assertFalse(reached_end) + + +class WaitForShieldTests(unittest.IsolatedAsyncioTestCase): + + async def test_zero_timeout(self): + # `asyncio.shield` creates a new task which wraps the passed in + # awaitable and shields it from cancellation so with timeout=0 + # the task returned by `asyncio.shield` aka shielded_task gets + # cancelled immediately and the task wrapped by it is scheduled + # to run. + + async def coro(): + await asyncio.sleep(0.01) + return 'done' + + task = asyncio.create_task(coro()) + with self.assertRaises(asyncio.TimeoutError): + shielded_task = asyncio.shield(task) + await asyncio.wait_for(shielded_task, timeout=0) + + # Task is running in background + self.assertFalse(task.done()) + self.assertFalse(task.cancelled()) + self.assertTrue(shielded_task.cancelled()) + + # Wait for the task to complete + await asyncio.sleep(0.1) + self.assertTrue(task.done()) + + + async def test_none_timeout(self): + # With timeout=None the timeout is disabled so it + # runs till completion. + async def coro(): + await asyncio.sleep(0.1) + return 'done' + + task = asyncio.create_task(coro()) + await asyncio.wait_for(asyncio.shield(task), timeout=None) + + self.assertTrue(task.done()) + self.assertEqual(await task, "done") + + async def test_shielded_timeout(self): + # shield prevents the task from being cancelled. + async def coro(): + await asyncio.sleep(0.1) + return 'done' + + task = asyncio.create_task(coro()) + with self.assertRaises(asyncio.TimeoutError): + await asyncio.wait_for(asyncio.shield(task), timeout=0.01) + + self.assertFalse(task.done()) + self.assertFalse(task.cancelled()) + self.assertEqual(await task, "done") + if __name__ == '__main__': unittest.main() diff --git a/Misc/NEWS.d/next/Library/2022-10-22-09-26-43.gh-issue-96764.Dh9Y5L.rst b/Misc/NEWS.d/next/Library/2022-10-22-09-26-43.gh-issue-96764.Dh9Y5L.rst new file mode 100644 index 00000000000000..a0174291cbc311 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2022-10-22-09-26-43.gh-issue-96764.Dh9Y5L.rst @@ -0,0 +1 @@ +:func:`asyncio.wait_for` now uses :func:`asyncio.timeout` as its underlying implementation. Patch by Kumar Aditya. pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy