Skip to main content

Command Palette

Search for a command to run...

Building Context Managers: enter, exit, and @contextmanager

The with statement isn't magic. It's just two method calls wrapped in a try/finally.

Published
9 min read
Building Context Managers: enter, exit, and @contextmanager
M
I'm Moussa, a Beninese Software Engineer based in Rwanda. I took what many considered an unlikely path, pivoting from Diplomacy and International Relations to Software Engineering. That career switch taught me something I carry into everything I do: the best learning happens when you're willing to start from zero. I work as a full-time software engineer and contribute to nonprofits supporting African students and sport talents in their education. I'm passionate about EdTech and AI and their potential to empower learners across Sub-Saharan Africa. Writing is how I process what I learn and give it back — every article here is something I wish I'd found when I was figuring it out myself. Fun facts: I'm a certified IJF judo instructor, I taught myself English through YouTube, and I'm always one rabbit hole away from picking up something new.

In Part 2, we used with open(...) as a black box: it works, the file closes, everyone's happy. In Part 3, we learned how generators pause and resume execution with yield. Now we put both together and open the black box.

This article is about building your own context managers from scratch, understanding exactly what Python calls and when, and choosing between two approaches depending on what you're building.


What Python Actually Does with with

When Python encounters a with statement, it doesn't do anything magical. It calls two methods in a specific order:

with open("notes.txt", "r") as f:
    content = f.read()

Python translates this into roughly:

manager = open("notes.txt", "r")
f = manager.__enter__()          # setup
try:
    content = f.read()
finally:
    manager.__exit__(None, None, None)  # teardown — always runs

That's the entire contract. Any object that has __enter__ and __exit__ methods is a context manager. Nothing more required.


__enter__: The Setup Method

__enter__ runs the moment execution enters the with block. Whatever it returns is what gets bound to the as name.

class Example:
    def __enter__(self):
        print("entering")
        return "hello"   # this is what `as x` receives

    def __exit__(self, exc_type, exc_value, traceback):
        print("exiting")
        return False

with Example() as x:
    print(x)             # "hello"

Output:

entering
hello
exiting

__enter__ can return anything: the object itself, a resource it manages, or nothing at all. For file objects, __enter__ returns self, the file object, which is why f in as f is the file you write to.


__exit__: The Teardown Method

__exit__ always runs when the block exits. It receives three arguments that describe how the block exited:

def __exit__(self, exc_type, exc_value, traceback):
    ...
Parameter What it contains
exc_type The exception class, e.g. ValueError
exc_value The exception instance
traceback The traceback object

If the block exited normally, that means no exception and all three are None.

The return value matters. Return True to suppress the exception. Return False (or None) to let it propagate:

def __exit__(self, exc_type, exc_value, traceback):
    return True   # swallows any exception
    return False  # re-raises any exception

Getting this wrong is a common mistake as we'll see it again shortly.


Building a Class-Based Context Manager

Let's build something real. A DatabaseConnection class that opens a SQLite connection on entry, commits on clean exit, and rolls back on error:

import sqlite3

class DatabaseConnection:
    def __init__(self, db_path):
        self.db_path = db_path
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_path)
        return self.connection           # caller gets the connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            self.connection.rollback()  # something went wrong — undo
        else:
            self.connection.commit()    # all good — save changes
        self.connection.close()
        return False                    # never suppress exceptions

Using it:

with DatabaseConnection("app.db") as conn:
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE IF NOT EXISTS users (name TEXT, age INT)")
    cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")
# commit and close happen automatically

# If an exception occurred inside the block:
try:
    with DatabaseConnection("app.db") as conn:
        cursor = conn.cursor()
        cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")
        raise ValueError("something went wrong")
except ValueError:
    pass
# rollback and close happened automatically — database unchanged

This pattern — commit on success, rollback on failure, always close — is exactly how Django's database transaction handling works internally.


A Common Mistake: The Wrong Return Value

Here's a bug that's easy to miss:

import time

class Timer:
    def __enter__(self):
        self.start = time.time()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.time() - self.start
        return self.elapsed    # ⚠️ this is a float — truthy!

self.elapsed is something like 0.234 — a truthy value. So __exit__ suppresses every exception that occurs inside the with block:

with Timer() as t:
    raise ValueError("something broke")

print("this prints — exception silently swallowed")  # 😬

Always be explicit:

def __exit__(self, exc_type, exc_value, traceback):
    self.elapsed = time.time() - self.start
    return False   # never suppress
💡
Please run the code to better understand it. It is crucial for the rest of the article.

The @contextmanager Approach

The class approach works well, but it has a structural awkwardness: the setup and teardown are split across two methods. You have to store intermediate state, like self.start in the Timer, as instance variables just to pass it between __enter__ and __exit__.

@contextmanager from contextlib solves this by letting you write a context manager as a single generator function. The yield is the dividing line:

from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()      # setup — before yield
    yield                    # with block runs here
    elapsed = time.time() - start  # teardown — after yield
    print(f"Elapsed: {elapsed:.4f}s")

The setup and teardown live in the same function, share the same local scope, and flow naturally top to bottom. No class needed, no instance variables.


How @contextmanager Works Internally

This is where Article 3 pays off. When Python sees yield inside a function, that function becomes a generator. @contextmanager wraps that generator and wires it to __enter__ and __exit__:

__enter__ is called
    → runs the generator up to yield
    → pauses there
    → returns the yielded value (bound to `as`)

your with block runs

__exit__ is called
    → resumes the generator after yield
    → runs the rest of the function (cleanup)

Simplified, this is what @contextmanager does internally:

class _GeneratorContextManager:
    def __init__(self, gen):
        self.gen = gen

    def __enter__(self):
        return next(self.gen)    # run up to yield, return yielded value

    def __exit__(self, exc_type, exc_value, traceback):
        try:
            next(self.gen)       # resume after yield — run cleanup
        except StopIteration:
            pass                 # generator finished normally
        return False

It's just a generator being driven by two next() calls. The yield really is a pause button.


Yielding a Value

Whatever you yield is what the caller receives via as:

@contextmanager
def database_connection(db_path):
    conn = sqlite3.connect(db_path)
    try:
        yield conn               # caller gets the connection
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

with database_connection("app.db") as conn:
    cursor = conn.cursor()
    cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")

Compare this to the class version from earlier. Same behavior, fewer lines, and the logic reads top to bottom instead of being split across two methods.


Handling Exceptions with @contextmanager

When an exception occurs inside the with block, it gets thrown into the generator at the yield point. Wrap the yield in try/except to handle it:

@contextmanager
def managed_operation(name):
    print(f"Starting {name}")
    try:
        yield
    except ValueError as e:
        print(f"Handled ValueError: {e}")
        # not re-raising — exception suppressed
    except Exception:
        print("Unhandled exception — re-raising")
        raise             # re-raise anything else
    finally:
        print(f"Cleaning up {name}")  # always runs


with managed_operation("task A"):
    print("Working...")

print("---")

with managed_operation("task B"):
    raise ValueError("something went wrong")

print("Execution continues — ValueError was suppressed")

Output:

Starting task A
Working...
Cleaning up task A
---
Starting task B
Handled ValueError: something went wrong
Cleaning up task B
Execution continues — ValueError was suppressed

The finally block is your guarantee that cleanup runs regardless of what happened. Use except only when you specifically want to handle or suppress a particular exception.


The One Rule: Yield Exactly Once

@contextmanager requires the generator to yield exactly one time. Not zero times. Not two times. Once.

# ❌ no yield — __enter__ will crash
@contextmanager
def broken():
    print("setup")
    # forgot yield

# ❌ two yields — __exit__ will crash
@contextmanager
def also_broken():
    yield "first"
    yield "second"   # not allowed

# ✅ exactly one yield
@contextmanager
def correct():
    print("setup")
    yield
    print("teardown")

If you yield zero times, __enter__ raises RuntimeError. If you yield more than once, __exit__ raises RuntimeError. The decorator enforces this strictly.


When to Use Each Approach

Both approaches produce objects that work identically with with. The choice is about which one fits the situation better.

Situation Use
Simple setup/teardown, no extra state @contextmanager
Complex logic with multiple methods Class with __enter__/__exit__
Need to subclass or extend the manager Class approach
Sharing state between enter and exit @contextmanager (use local variables)
Most real-world cases @contextmanager

In practice, @contextmanager covers the vast majority of use cases. The class approach is mainly useful when the context manager is complex enough to warrant its own type, or when you need to build a hierarchy of context managers.


A Quick Side-by-Side

The same context manager written both ways:

# Class-based
class ManagedFile:
    def __init__(self, path):
        self.path = path

    def __enter__(self):
        self.file = open(self.path, "w", encoding="utf-8")
        return self.file

    def __exit__(self, exc_type, exc_value, traceback):
        self.file.close()
        return False


# Function-based
from contextlib import contextmanager

@contextmanager
def managed_file(path):
    file = open(path, "w", encoding="utf-8")
    try:
        yield file
    finally:
        file.close()

Both work identically. The function-based version is 5 lines shorter and reads more naturally. Unless you have a specific reason to use a class, @contextmanager is the default choice.


The Mental Model to Take Away

The with statement calls __enter__ on the way in and __exit__ on the way out — always. That's the entire contract.

@contextmanager wires a generator function to that contract automatically. Everything before yield maps to __enter__. Everything after yield maps to __exit__. The generator's pause mechanism is exactly what makes this work, and now you understand why.

Three things to carry forward:

  • __exit__ receives exception info — check exc_type is not None to detect errors

  • Return False from __exit__ unless you specifically intend to suppress an exception

  • With @contextmanager, always wrap yield in try/finally when there's cleanup to do


Acronyms Used in This Article

  • PEP — Python Enhancement Proposal. PEP 343 introduced the with statement and the context manager protocol.

This is Part 4 of 5 in the Python Context Managers series. Next up: Part 5 — Context Managers in the Wild

Python Context Managers: A Deep Dive

Part 1 of 4

A five-part series on Python context managers, built for developers coming from JavaScript. We start with the resource leak problem, work through the with statement, open(), generator functions, and building your own context managers, then finish with real-world usage across popular frameworks.

Up next

Generator Functions in Python: How to Pause a Function Mid-Execution

A function that remembers where it left off changes everything.