diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index b4f1eed91a9321..ef5c0f9669abbb 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 @@ -442,62 +443,12 @@ async def wait_for(fut, timeout): This function is a coroutine. """ - loop = events.get_running_loop() - - if timeout is None: - return await fut - - if timeout <= 0: - fut = ensure_future(fut, loop=loop) - - if fut.done(): - return fut.result() - - await _cancel_and_wait(fut, loop=loop) - 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) + async def inner(): + async with timeouts.timeout(timeout): + return await fut - 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() + return await create_task(inner()) async def _wait(fs, timeout, return_when, loop):