Writing Your Own Python Context Manager, With a Class or contextlib

Write your own context manager when something must be undone no matter how the block ends: a lock released, a setting restored, a temporary file deleted, a timer stopped. with guarantees the clean-up runs even if the code inside raises an error. If nothing needs undoing, you don't need one.

What with does

with manager as value:
    do_work(value)

Python calls manager.__enter__(), and whatever it returns goes into value. Then the block runs. When the block ends, by finishing, return, break or an exception, Python calls manager.__exit__(...). That last part is the whole point. You can't forget the clean-up, and an error can't skip it.

A real job: change folder, then always change back

Some tools must run inside a certain folder. If you os.chdir() there and an error happens, your program stays in the wrong folder, and every relative path after that points somewhere else.

The class way

import os


class WorkingDirectory:
    def __init__(self, path):
        self.path = path
        self.previous = None

    def __enter__(self):
        self.previous = os.getcwd()
        os.chdir(self.path)
        return self.path

    def __exit__(self, exc_type, exc_value, traceback):
        os.chdir(self.previous)
        return False


os.makedirs("build", exist_ok=True)
start = os.getcwd()

with WorkingDirectory("build") as folder:
    print("inside:", os.path.basename(os.getcwd()))

try:
    with WorkingDirectory("build"):
        raise ValueError("the build failed")
except ValueError as error:
    print("error:", error)

print("back where we started:", os.getcwd() == start)

inside: build
error: the build failed
back where we started: True

The error happened inside the block, and the folder was still restored. __exit__ gets three arguments: the exception type, the exception, and the traceback. All three are None when the block ended normally.

Should it swallow the error?

__exit__ returning False (or None) means "I cleaned up, now let the error continue". Returning True means "I handled it, pretend it never happened".

class IgnoreEverything:
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return True


with IgnoreEverything():
    total = 100 / 0
print("still running, and total is never set")

still running, and total is never set

That's the bare except: problem in a nicer costume. Return True only for a specific exception type that you really handle, and let the rest through. Most context managers should return False. The standard library has one that suppresses on purpose, contextlib.suppress(FileNotFoundError), and you have to name the error.

The contextlib way

For most cases, write a generator function and add @contextmanager:

import time
from contextlib import contextmanager


@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed:.1f}s")


with timer("sleep"):
    time.sleep(0.2)

try:
    with timer("failed step"):
        time.sleep(0.1)
        raise RuntimeError("step failed")
except RuntimeError as error:
    print("error:", error)

sleep: 0.2s
failed step: 0.1s
error: step failed

Everything before yield is __enter__. The value you yield goes to as. Everything after is __exit__.

Giving a value to as

import os
import shutil
import tempfile
from contextlib import contextmanager
from pathlib import Path


@contextmanager
def scratch_folder():
    path = Path(tempfile.mkdtemp())
    try:
        yield path
    finally:
        shutil.rmtree(path)


with scratch_folder() as folder:
    (folder / "draft.txt").write_text("temporary", encoding="utf-8")
    print("files inside:", os.listdir(folder))
    kept = folder

print("deleted after:", not kept.exists())

files inside: ['draft.txt']
deleted after: True

Again, the standard library already has this as tempfile.TemporaryDirectory(). When you need a context manager, check contextlib and tempfile first. Writing your own is for your own resources: a connection to your system, a setting in your app, a lock in your code.

Class or contextlib?

• **@contextmanager:** short set-up and clean-up, no extra methods. Most cases.
• A class: the object needs other methods or state you read later (like a timer you can ask for elapsed after the block), or it needs to be reused many times with the same settings.

Common mistakes

• **No try/finally around yield.** The clean-up skips on errors.
• **yield more than once.** @contextmanager raises RuntimeError: generator didn't stop.
• **Returning True from __exit__ by accident.** Every error inside disappears.
• Writing one that the standard library already has. open, chdir, suppress, TemporaryDirectory, redirect_stdout.

What happens if you skip it

You write start, work, stop as three lines. It works until work raises an error. Then the lock is never released, the setting stays changed for the rest of the program, or the temp folder stays on disk forever. One error turns into many, far from where it happened.

Which clean-up in your code today only runs if nothing goes wrong?