Python Context Managers: Clean Up Even When Things Fail

A context manager is anything you can use with with. Its job is to set something up and guarantee the clean-up runs afterwards, even if the code in the middle fails. Files are the example everyone knows. Locks, database transactions, temporary folders and "change this setting, then put it back" are the same idea.

What with actually does

This:

with open("data.txt", encoding="utf-8") as f:
    process(f)

is roughly this:

manager = open("data.txt", encoding="utf-8")
f = manager.__enter__()
try:
    process(f)
finally:
    manager.__exit__(...)   # always runs, error or not

The try/finally is the whole point. You could write it by hand every time. People don't, or they forget once, which is why with exists.

Writing one as a class

A timer that reports how long a block took, including when the block fails:

import time


class Timer:
    def __init__(self, label):
        self.label = label

    def __enter__(self):
        self.start = time.perf_counter()
        return self                      # this is what "as t" receives

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.perf_counter() - self.start
        status = "failed" if exc_type else "ok"
        print(f"{self.label}: {status} after {self.elapsed:.1f}s")
        return False                     # don't hide the error


with Timer("import") as t:
    time.sleep(0.3)

try:
    with Timer("export"):
        time.sleep(0.1)
        raise ValueError("disk full")
except ValueError as error:
    print("caught:", error)

import: ok after 0.3s
export: failed after 0.1s
caught: disk full

__exit__ gets three arguments describing the error, or three None values if the block finished normally. What it returns decides what happens to the error.

The shorter way: @contextmanager

For most cases, a generator function with contextlib.contextmanager is less code. Everything before yield is the setup, the yielded value is what as receives, and everything after is the clean-up:

import os
from contextlib import contextmanager


@contextmanager
def working_directory(path):
    previous = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(previous)          # runs even if the block fails


os.makedirs("exports", exist_ok=True)
start = os.getcwd()
with working_directory("exports"):
    print(os.path.basename(os.getcwd()))
print(os.getcwd() == start)

exports
True

The try/finally around yield is required. Without it, an error inside the with block is raised at the yield, and the clean-up line after it never runs. It is the most common mistake with @contextmanager.

A real-world example: write all or nothing

A file that is either fully written or not touched at all. If anything fails while writing, the original file stays as it was:

import os
from contextlib import contextmanager


@contextmanager
def atomic_write(path, encoding="utf-8"):
    tmp = path + ".tmp"
    f = open(tmp, "w", encoding=encoding)
    try:
        yield f
        f.close()
        os.replace(tmp, path)       # only reached if the block succeeded
    except BaseException:
        f.close()
        os.remove(tmp)
        raise


with atomic_write("prices.csv") as f:
    f.write("sku,price\nKB-01,1200\n")

try:
    with atomic_write("prices.csv") as f:
        f.write("sku,price\n")
        raise RuntimeError("source API went down halfway")
except RuntimeError as error:
    print("failed:", error)

with open("prices.csv", encoding="utf-8") as f:
    print(f.read(), end="")
print(os.path.exists("prices.csv.tmp"))

failed: source API went down halfway
sku,price
KB-01,1200
False

The second write failed halfway, and prices.csv still has the complete first version. No half-written file, no leftover temp file. Every script that writes the same file would need this logic. As a context manager, it is one with line.

Useful ones that already exist

import os
import tempfile
from contextlib import suppress

with tempfile.TemporaryDirectory() as tmp:
    path = os.path.join(tmp, "scratch.txt")
    with open(path, "w", encoding="utf-8") as f:
        f.write("temporary")
    print(os.path.exists(path))
print(os.path.exists(path))          # the folder and its contents are gone

with suppress(FileNotFoundError):
    os.remove("does-not-exist.txt")  # no error if it isn't there
print("continued")

True
False
continued

Also worth knowing: threading.Lock() works with with and always releases the lock, and database libraries usually offer a transaction as a context manager that commits on success and rolls back on error. Check the library's documentation for the exact behaviour before relying on it.

Common mistakes

• No try/finally around yield in a @contextmanager function. The clean-up silently doesn't run on errors.
• Using suppress for errors you don't understand. It is for one specific, expected error. suppress(Exception) is a bare except: pass with better manners.
• Using the object after the block. After with open(...) as f: ends, f still exists but is closed. Reading from it raises ValueError: I/O operation on closed file.

What happens if you skip this step

Clean-up code runs on the happy path and gets skipped on the error path. Locks stay locked, temp files pile up, connections leak until the database refuses new ones. These bugs only show up after an error, under load, often days later, which is why they are hard to trace back.

Your script changes a shared setting at the start and restores it at the end. What state is that setting in after the script crashes in the middle?