Handle Errors and Nesting in a Context Manager

In the last exercise, Borrowed checked out a connection in __enter__ and gave it back in __exit__. In real code, a helper function might enter an object that its caller has already entered:

conn = PooledResource(pool)
with conn:
    with conn:  # a helper enters the same object again
        run_query(conn.name)
    run_query(conn.name)  # still inside the outer block

If __enter__ checks out a connection every time, the inner with takes a second one. Then the inner exit gives a connection back while the outer block still needs it. This article sets out the rules with follows, and the exercise after it fixes a context manager that breaks them.

What with promises

• with calls __enter__ first. If __enter__ raises, the block does not run and with does not call __exit__. So if __enter__ fails halfway, it must undo its own work.
• Once __enter__ returns, with always calls __exit__. That happens when the block finishes, when it raises, and when it leaves through return, break or continue.
• If __exit__ returns a true value, Python hides the block's exception. Return None to let the exception through.
• If __exit__ itself raises, its exception replaces the one from the block.

with does not free the object. After the block, the name and the object both still exist:

import tempfile

with tempfile.TemporaryFile(mode="w+t") as handle:
    handle.write("example")

print(handle.closed)  # True: the file is closed, but handle still exists

A pooled-connection wrapper works the same way. It gives its connection back at the end of the block, and the wrapper object lives on until nothing points at it.

Make an object safe to enter twice

An object that you can enter again while it is already entered is called reentrant. To make one, keep a depth count, the number of with blocks that have entered the object and not yet exited. Check out a connection only when the count goes from 0 to 1, and give it back only when the count returns to 0. Inner entries use the connection the outer entry holds.

Let mypy catch return True

When __enter__ returns the object itself, annotate it with Self from typing. A subclass then gets its own type back from with. When __exit__ should never hide an error, annotate it -> None. mypy then rejects an accidental return True:

from types import TracebackType
from typing import Self


class Quiet:
    def __enter__(self) -> Self:
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        return True  # mypy: No return value expected

The annotation changes nothing when the program runs. It gives mypy something to check.

Going further

The with statement reference gives the full rules. contextlib.ExitStack enters a number of context managers that you know only at run time, and exits them all in reverse order.