Write Async Context Managers That Clean Up

Each request to your service opens a session to another API. You must close the session when the request ends, whether the handler returns, raises or is cancelled. An async context manager does this job. You use it with async with, and it can await while it sets up and while it cleans up. The async contextlib article covers the standard helpers. This article covers the mistakes that matter in concurrent code.

The shortest way to write one is contextlib.asynccontextmanager:

import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager


class Session:
    async def close(self) -> None:
        await asyncio.sleep(0)  # stands in for closing a connection
        print("session closed")


@asynccontextmanager
async def open_session() -> AsyncIterator[Session]:
    session = Session()
    try:
        yield session
    finally:
        await session.close()


async def main() -> None:
    try:
        async with open_session():
            raise ValueError("handler failed")
    except ValueError as exc:
        print(f"caller still sees: {exc}")


asyncio.run(main())

session closed
caller still sees: handler failed

The code before yield sets up, and the value you yield is what as binds. If the body raises, the error appears at the yield, the finally runs, and the error propagates to the caller.

Re-raise after you log

If you catch the error at the yield and let the function finish, the manager suppresses the error. It stops there, and the caller continues as if the body worked. A log line is easy to mistake for handling:

try:
    yield session
except Exception:
    log.exception("request failed")
    raise  # without this line, the caller never sees the error
finally:
    await session.close()

Write the class form with accurate types

When you write the class yourself, __aenter__ returns the value that as binds, and __aexit__ receives the exception, if there is one:

from types import TracebackType
from typing import Self


class Session:
    async def __aenter__(self) -> Self:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        await self.close()

    async def close(self) -> None: ...

Annotate __aenter__ with Self when it returns self, so a subclass keeps its own type. If it returns something else, such as a connection, use that type.

The return value of __aexit__ decides what happens to the exception. None or False lets it propagate to the caller. True suppresses it, even a CancelledError, so a cancelled task keeps running. Annotate -> None when the manager should always let errors through. The type checker then rejects a stray return True.

Clean up a half-finished setup

If __aenter__ raises, Python skips __aexit__. Say __aenter__ opens two connections and the second one fails. Close the first one before the error leaves __aenter__. contextlib.AsyncExitStack can track what you have opened so far.

When a task is cancelled inside an async with, __aexit__ or the generator's finally runs with the CancelledError. Cleanup can still await. If it can hang, give it a limit with asyncio.timeout, as in the two previous articles.

Close async generators you stop early

When you break out of an async for over an async generator, the generator's finally does not run yet. It runs later, at a time you do not control. Wrap the generator in contextlib.aclosing to close it as the block ends:

async with aclosing(read_lines(path)) as lines:
    async for line in lines:
        if line == "END":
            break

The async iteration article covers async generators in more depth.

See the contextlib documentation.