[FIXED] Why is __aexit__ not fully executed when it has await inside?

Issue

This is the simplified version of my code:

main is a coroutine which stops after the second iteration.
get_numbers is an async generator which yields numbers but within an async context manager.

import asyncio


class MyContextManager:
    async def __aenter__(self):
        print("Enter to the Context Manager...")
        return self

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        print(exc_type)
        print("Exit from the Context Manager...")
        await asyncio.sleep(1)
        print("This line is not executed")  # <-------------------
        await asyncio.sleep(1)


async def get_numbers():
    async with MyContextManager():
        for i in range(30):
            yield i


async def main():
    async for i in get_numbers():
        print(i)
        if i == 1:
            break


asyncio.run(main())

And the output is:

Enter to the Context Manager...
0
1
<class 'asyncio.exceptions.CancelledError'>
Exit from the Context Manager...

I have two questions actually:

1. From my understanding, AsyncIO schedules a Task to be called soon in the next cycle of the event loop and gives __aexit__ a chance to execute. But the line print("This line is not executed") is not executed. Why is that? Is it correct to assume that if we have an await statement inside the __aexit__, the code after that line is not going to execute at all and we shouldn't rely on that for cleaning?

2. Output of the help() on async generators shows that:

 |  aclose(...)
 |      aclose() -> raise GeneratorExit inside generator.

so why I get <class 'asyncio.exceptions.CancelledError'> exception inside the __aexit__ ?

_{* I'm using Python 3.10.4}

Solution

This is not specific to __aexit__ but to all async code: When an event loop shuts down it must decide between cancelling remaining tasks or preserving them. In the interest of cleanup, most frameworks prefer cancellation instead of relying on the programmer to clean up preserved tasks later on.

This kind of shutdown cleanup is a separate mechanism from the graceful unrolling of functions, contexts and similar on the call stack during normal execution. A context manager that must also clean up during cancellation must be specifically prepared for it. Still, in many cases it is fine not to be prepared for this since many resources fail safe by themselves.

---

In contemporary event loop frameworks there are usually three levels of cleanup:

• Unrolling: The __aexit__ is called when the scope ends and might receive an exception that triggered the unrolling as an argument. Cleanup is expected to be delayed as long as necessary. This is comparable to __exit__ running synchronous code.
• Cancellation: The __aexit__ may receive a CancelledError¹ as an argument or as an exception on any await/async for/async with. Cleanup may delay this, but is expected to proceed as fast as possible. This is comparable to KeyboardInterrupt cancelling synchronous code.
• Closing: The __aexit__ may receive a GeneratorExit as an argument or as an exception on any await/async for/async with. Cleanup must proceed as fast as possible. This is comparable to GeneratorExit closing a synchronous generator.

To handle cancellation/closing, any async code – be it in __aexit__ or elsewhere – must expect to handle CancelledError or GeneratorExit. While the former may be delayed or suppressed, the latter should be dealt with immediately and synchronously².

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        print("Exit from the Context Manager...")
        try:
            await asyncio.sleep(1)
        except GeneratorExit:
            print("Exit stage left NOW")
        except asyncio.CancelledError:
            print("Got cancelled, just cleaning up a few things...")
            await asyncio.sleep(0.5)
        else:
            print("Nothing to see here, taking my time on the way out")
            await asyncio.sleep(1)

Note: It is generally not possible to exhaustively handle these cases. Different forms of cleanup may interrupt one another, such as unrolling being cancelled and then closed. Handling cleanup is only possible on a best effort basis; robust cleanup is achieved by fail safety, for example via transactions, instead of explicit cleanup.

---

Cleanup of asynchronous generators in specific is a tricky case since they can be cleaned up by all cases at once: Unrolling as the generator finishes, cancellation as the owning task is destroyed or closing as the generator is garbage collected. The order at which the cleanup signals arrive is implementation dependent.

The proper way to address this is not to rely on implicit cleanup in the first place. Instead, every coroutine should make sure that all its resources are closed before exiting.

async def main():
    async_iter = get_numbers()
    async for i in async_iter:
        print(i)
        if i == 1:
            break
    await async_iter.aclose()  # wait for generator to close before exiting

In recent versions, this pattern is codified via the aclosing context manager.

---

¹The name and/or identity of this exception may vary.

²While it is possible to await asynchronous things during GeneratorExit, they may not yield to the event loop. A synchronous interface is advantageous to enforce this.

Answered By - MisterMiyagi