Python asyncio.wait_for()

Summary: in this tutorial, you’ll learn how to use the asyncio.wait_for() function to wait for a coroutine to complete with a timeout.

Introduction to the Python asyncio.wait_for() function

In the previous tutorial, you learned how to cancel a task that is in progress by using the cancel() method of the Task object.

To wait for a task to complete with a timeout, you can use the asyncio.wait_for() function. The asyncio.wait_for() function waits for a single task to be completed with a timeout.

When a timeout occurs, the asyncio.wait_for() function cancels the task and raises the TimeoutError exception. Otherwise, it returns the result of the task. For example:

import asyncio
from asyncio.exceptions import TimeoutError


async def call_api(message, result=1000, delay=3):
    print(message)
    await asyncio.sleep(delay)
    return result


async def main():
    task = asyncio.create_task(
        call_api('Calling API...', result=2000, delay=5)
    )

    MAX_TIMEOUT = 3
    try:
        await asyncio.wait_for(task, timeout=MAX_TIMEOUT)
    except TimeoutError:
        print('The task was cancelled due to a timeout')

asyncio.run(main())
Code language: JavaScript (javascript)

Output:

Calling API...
The task was cancelled due to a timeout

How it works

First, define a call_api() coroutine that takes 3 seconds to complete by default:

async def call_api(message, result=1000, delay=3):
    print(message)
    await asyncio.sleep(delay)
    return resultCode language: JavaScript (javascript)

Second, create a task that wraps the call_api coroutine and takes 5 seconds to complete:

task = asyncio.create_task(
    call_api('Calling API...', result=2000, delay=5)
)Code language: JavaScript (javascript)

Third, use the asyncio.wait_for() function to wait for the task to complete with a timeout of 3 seconds. Since the task takes 5 seconds to complete, a timeout will occur and a TimeoutError will be raised:

MAX_TIMEOUT = 3
try:
    await asyncio.wait_for(task, timeout=MAX_TIMEOUT)
except TimeoutError:
    print('The task was cancelled due to a timeout')Code language: PHP (php)

Shielding a task from cancellation

Sometimes, you may want to inform users that a task is taking longer than expected after a certain amount of time but not cancel the task when a timeout is exceeded.

To do that, you can wrap the task with the asyncio.shield() function. The asyncio.shield() prevents the cancellation of a task. For example:

import asyncio
from asyncio.exceptions import TimeoutError


async def call_api(message, result=1000, delay=3):
    print(message)
    await asyncio.sleep(delay)
    return result


async def main():
    task = asyncio.create_task(
        call_api('Calling API...', result=2000, delay=5)
    )

    MAX_TIMEOUT = 3
    try:
        await asyncio.wait_for(asyncio.shield(task), timeout=MAX_TIMEOUT)
    except TimeoutError:
        print('The task took more than expected and will complete soon.')
        result = await task
        print(result)

asyncio.run(main())
Code language: JavaScript (javascript)

Output:

Calling API...
The task took more than expected and will complete soon.
2000

In this example, the task takes 5 seconds to complete. Once the timeout is 3 seconds, the TimeoutEror exception is raised. However, the task is not canceled due to the asyncio.shield() function.

In the exception handling section, we await for the task to be completed and print out the result.

Summary

• Use asyncio.wait_for() function to wait for a task with a timeout.
• Use asyncio.shield() function to prevent the cancellation of a task after a timeout.