Context Managers

Many jobs come in pairs: open a file and close it, lock a door and unlock it, start a database transaction and either save or undo it. Forgetting the second half causes real bugs, such as files that never get written or data left half-changed after an error. A context manager is an object that does the "set up" part when a block starts and the "clean up" part when it ends, however it ends.

You have used one already: with open(...) as f:. In this lesson you will see what with really does, write your own context managers in two ways, and use the ready-made helpers in the contextlib module.

What the with statement guarantees

When a with block finishes, the clean-up runs. That is true if the block ends normally, if it uses return or break, and if an exception is raised inside it.

from pathlib import Path

with open("attendance.txt", "w", encoding="utf-8") as f:
    f.write("Aarav,present\n")
    f.write("Priya,absent\n")
print("Closed after the block?", f.closed)

try:
    with open("attendance.txt", "a", encoding="utf-8") as f:
        f.write("Kabir,present\n")
        raise RuntimeError("power cut!")
except RuntimeError as error:
    print("Error:", error, "| file closed?", f.closed)

print(Path("attendance.txt").read_text(encoding="utf-8"))

Closed after the block? True
Error: power cut! | file closed? True
Aarav,present
Priya,absent
Kabir,present

Even when the error struck, the file was closed and Kabir's line was saved. Without with, you would need this longer try/finally form, which you will still see in older code:

f = open("attendance.txt", "a", encoding="utf-8")
try:
    f.write("Meera,present\n")
finally:
    f.close()          # runs whether or not an error happened

How with works: __enter__ and __exit__

Any object with two special methods can be used in a with statement. Python calls __enter__() at the start, and whatever it returns is bound to the name after as. At the end, Python calls __exit__(exc_type, exc_value, traceback). If the block raised an exception, those three arguments describe it; otherwise all three are None. You saw a simple version in Special (Dunder) Methods.

The arguments to __exit__ are what make context managers powerful. Here is a manager that lets you change a shop's stock, and puts everything back if anything goes wrong halfway:

class StockUpdate:
    def __init__(self, stock):
        self.stock = stock

    def __enter__(self):
        self._backup = dict(self.stock)       # remember the old state
        return self.stock

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:              # an error happened inside
            self.stock.clear()
            self.stock.update(self._backup)
            print(f"Rolled back because of {exc_type.__name__}: {exc_value}")
        return False                          # let the error continue

stock = {"rice": 50, "dal": 30}

with StockUpdate(stock) as s:
    s["rice"] -= 5
print(stock)

try:
    with StockUpdate(stock) as s:
        s["rice"] -= 10
        s["dal"] -= 40
        if s["dal"] < 0:
            raise ValueError("not enough dal")
except ValueError:
    print("Order failed")
print(stock)

{'rice': 45, 'dal': 30}
Rolled back because of ValueError: not enough dal
Order failed
{'rice': 45, 'dal': 30}

The second order took rice and then failed on dal. Without the rollback, the rice would be gone for an order that never happened. This is exactly how database transactions behave, as you will see in SQLite Databases.

The return value of __exit__ matters. Returning False (or nothing) lets the exception carry on to the caller after the clean-up. Returning True swallows the exception. Only do that on purpose, because hidden errors are hard to debug.

The easy way: @contextmanager

Writing a class is a lot of typing for a small job. The contextlib.contextmanager decorator turns a generator function (a function with yield, see Iterators & Generators) into a context manager. Code before yield is the set-up, the yielded value goes to as, and code after yield is the clean-up. Wrap the yield in try/finally so the clean-up runs even after an error.

from contextlib import contextmanager

@contextmanager
def exam_hall(room):
    print(f"Unlock {room}, switch on fans")
    try:
        yield f"{room} register"
    finally:
        print(f"Collect papers, lock {room}")

with exam_hall("Room 12") as register:
    print("Writing exam, signing", register)

try:
    with exam_hall("Room 14"):
        raise RuntimeError("fire alarm")
except RuntimeError as error:
    print("Stopped:", error)

Unlock Room 12, switch on fans
Writing exam, signing Room 12 register
Collect papers, lock Room 12
Unlock Room 14, switch on fans
Collect papers, lock Room 14
Stopped: fire alarm

Room 14 was still locked up after the fire alarm, because the finally part ran. Forgetting try/finally around yield is the most common bug with this decorator: the clean-up would be skipped whenever the block fails.

Handy helpers in contextlib

The contextlib module includes ready-made context managers for everyday jobs:

• suppress(SomeError) ignores one specific exception type inside the block. It is a short, clear replacement for try/except/pass.
• redirect_stdout(buffer) sends everything that print() writes into another file-like object, which is handy for capturing output in tests.
• closing(thing) calls thing.close() at the end, for objects that have close() but no with support.
• nullcontext() does nothing; it is a placeholder when code sometimes needs a real manager and sometimes not.
• chdir(path) (Python 3.11 and later) changes the current folder and changes it back afterwards.

import io
from contextlib import suppress, redirect_stdout
from pathlib import Path

with suppress(FileNotFoundError):
    Path("old-report.txt").unlink()      # fine if the file is not there
print("Clean-up done")

buffer = io.StringIO()
with redirect_stdout(buffer):
    print("Order 1042 packed")
    print("Order 1043 packed")
captured = buffer.getvalue()
print("Captured", captured.count("\n"), "lines:", captured.split())

Clean-up done
Captured 2 lines: ['Order', '1042', 'packed', 'Order', '1043', 'packed']

Several resources at once

You can open several context managers in one with statement. Since Python 3.10 you can wrap them in brackets and put each on its own line:

with (
    open("marks.csv", encoding="utf-8") as source,
    open("toppers.csv", "w", encoding="utf-8") as target,
):
    for line in source:
        if int(line.split(",")[1]) >= 90:
            target.write(line)

When you do not know how many resources you need until the program runs, use contextlib.ExitStack. You add context managers to the stack one by one, and when the with block ends it closes all of them, in reverse order, even if one of them failed.

from contextlib import ExitStack
from pathlib import Path

cities = ["Rohtak", "Jaipur", "Pune"]
for city in cities:
    Path(f"{city}.txt").write_text(f"Orders from {city}: {len(city) * 3}\n", encoding="utf-8")

with ExitStack() as stack:
    files = [stack.enter_context(open(f"{city}.txt", encoding="utf-8")) for city in cities]
    stack.callback(print, "Closing all files")      # runs when the block ends
    for f in files:
        print(f.readline().strip())
    print("Any closed yet?", any(f.closed for f in files))

print("All closed now?", all(f.closed for f in files))

Orders from Rohtak: 18
Orders from Jaipur: 18
Orders from Pune: 12
Any closed yet? False
Closing all files
All closed now? True

enter_context() enters a manager and registers its clean-up. callback() registers any function to run at the end. Clean-ups run in reverse order, so the callback, which was added last, ran first, and then the three files were closed.

Common mistakes

• No try/finally around yield. The clean-up after yield is skipped when the block raises an error.
• Returning True from __exit__ by accident. Every error inside the block silently disappears.
• Using the resource after the block. f.read() after the with ends raises ValueError: I/O operation on closed file.
• Suppressing broad exceptions. suppress(Exception) hides real bugs; name the exact error you expect.

Key points

• A context manager pairs set-up with clean-up, and with guarantees the clean-up runs.
• Python calls __enter__() at the start and __exit__(exc_type, exc_value, traceback) at the end.
• __exit__ returning True swallows the exception; returning False lets it continue.
• @contextmanager builds a manager from a generator; put yield inside try/finally.
• suppress, redirect_stdout, closing, nullcontext and chdir cover common jobs.
• Use bracketed multi-manager with for a fixed set of resources and ExitStack for a variable number.

from contextlib import contextmanager

@contextmanager
def timer_note(task):
    print(f"Starting {task}")
    try:
        yield
    except Exception as error:
        print(f"{task} failed: {error}")
        raise                      # let the error continue
    else:
        print(f"{task} finished")

with timer_note("packing order"):
    print("Packing 3 items")

try:
    with timer_note("address check"):
        raise ValueError("bad pincode")
except ValueError:
    print("Caller handled the error")

Starting packing order
Packing 3 items
packing order finished
Starting address check
address check failed: bad pincode
Caller handled the error