from . import events
from . import exceptions
from . import futures
+from . import timeouts
from .coroutines import _is_coroutine
# Helper to generate new task names
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().
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)
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()
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()
projects / thirdparty / Python / cpython.git / commitdiff
