Asyncio.wait_for() function in Python

If you're looking to build efficient and fast I/O-bound applications, then the asyncio module in Python is a great tool to have in your arsenal. The wait_for() function is a crucial component of the asyncio module, allowing you to manage long-running functions in your application and make them awaitable. In this tutorial, we'll take a deep dive into what the wait_for() function is, how it works, and examples of how to use it in your own Python code. Let's get started!

Python asyncio.wait_for() Function :

• The  asyncio.wait_for()  function is used to wait for a routine to complete or raise an exception for a maximum amount of time.
• The function returns the result of the co-routine if it completes within the specified time frame, or raises a      TimeoutError  if it takes longer than the specified time.
• The function can also be used with multiple co-routines using the  asyncio.gather()  function.

Different use cases of Python asyncio.wait_for() :

•  asyncio.wait_for()  function can be used for :
  • Waiting for a Single Co-routine
  • Shield a Task from Cancellation

Waiting for a Single Coroutine Using Python asyncio.wait_for() :

• Here's an example of using  asyncio.wait_for()  to wait for a single co-routine to complete within a certain time frame:
• Code Sample : 

import asyncio

async def long_operation():
   print('Long operation started')
   await asyncio.sleep(5)
   print('Long operation completed')
   return 'Long operation result'

async def main():
   print('Main started')
   try:
       result = await asyncio.wait_for(long_operation(), timeout=3)
   except asyncio.TimeoutError:
       print('Timeout occurred')
   else:
       print('Result:', result)
   print('Main completed')


asyncio.run(main())

Explanation :

• In the above example, we define a co-routine long_operation() that waits for five seconds before returning.
• We also define a  main()  co-routine that calls  asyncio.wait_for()  with  long_operation()  as the co-routine argument and  timeout=3  as the timeout value.
• If  long_operation()  completes within three seconds, the result is printed to the console. Otherwise, a timeout error is raised and caught by the try/except block.

How to Shield a Task from Cancellation with Python asyncio.wait_for() :

• We can use  asyncio.shield()  to shield a co-routine from being cancelled while it's waiting for another co-routine to complete using  asyncio.wait_for() .
• Here's an example of how to shield a long-running task that is awaiting the completion of another co-routine using  asyncio.wait_for() :
• Code Sample : 

import asyncio

async def long_operation():
   print('Long operation started')
   await asyncio.sleep(5)
   print('Long operation completed')
   return 'Long operation result'

async def short_operation():
   print('Short operation started')
   await asyncio.sleep(2)
   print('Short operation completed')
   return 'Short operation result'

async def main():
   print('Main started')
   try:
       short_op_task = asyncio.create_task(short_operation())
       long_op_task = asyncio.create_task(asyncio.wait_for(long_operation(), timeout=4))
       
       result = await long_op_task
   except asyncio.TimeoutError:
       print('Timeout occurred during long operation')
   except Exception as exc:
       print(f'Exception occurred during long operation: {exc!r}')
   else:
       print('Long operation result:', result)
   await short_op_task
   print('Short operation result:', short_op_task.result())
   print('Main completed')
asyncio.run(main())

Explanation :

• In the above example, we define two co-routines,  long_operation()  and  short_operation() .
• We use  asyncio.create_task()  to create tasks for both co-routines.
• use  asyncio.wait_for()  to wait for long_operation() to complete within a timeout of four seconds.
• We shield  long_operation()  from being cancelled by wrapping it in  asyncio.shield() .
• If  long_operation()  completes within the timeout, its result is printed to the console.
• Whether or not  long_operation()  completes within the timeout,  short_operation()  is still executed and its result is printed to the console.