aboutsummaryrefslogtreecommitdiffstatshomepage
path: root/Lib/asyncio/tasks.py
diff options
context:
space:
mode:
Diffstat (limited to 'Lib/asyncio/tasks.py')
-rw-r--r--Lib/asyncio/tasks.py77
1 files changed, 29 insertions, 48 deletions
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index e78719de216..a2e06d5ef72 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -24,6 +24,7 @@ from . import coroutines
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 @@ async def _wait(fs, timeout, return_when, loop):
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)