Python Custom Context Managers

In the last lesson you learned about Closures. An inner function remembers values from the function around it. Now we use that same idea of remembering and cleaning up, but for whole blocks of code. You have already used with open(...) to work with files. The file closes itself when the block ends. It closes even if something goes wrong. The nice part is you can give your own objects that same behavior. That is what a custom context manager is.

🤔 Why Build Your Own Context Manager?

Picture this. You open something. You use it. Then you must close it. Maybe it is a database connection. Maybe a network socket. Maybe a lock or a timer. The pain is the closing step. People forget it. Or an error jumps over it and the cleanup never runs.

Here is the messy version most people write first.

import time

start = time.perf_counter()
total = sum(range(1_000_000))   # the work we want to measure
end = time.perf_counter()
print(f"Took {end - start:.4f} seconds")

That works. But the setup and the cleanup are loose lines floating around your real work. If an error happens between them, the final print never runs. A context manager bundles the “before” step and the “after” step into one object. And Python promises the “after” step always runs.

The deal is simple.

• Setup happens when the with block starts.
• Your code runs inside the block.
• Cleanup happens when the block ends. It runs even if an error was raised.

🧩 How the with Statement Works

Think of a context manager like checking into a hotel room. When you arrive, the front desk hands you a key. That is the setup. You stay and use the room. When you leave, they take the key and clean the room. That is the cleanup. You do not have to remember to clean it yourself. It just happens on the way out.

For Python to do this with your object, the object needs two special methods.

When you write with thing as x:, Python calls thing.__enter__() and puts its return value into x. When the block finishes, Python calls thing.__exit__(...). This pair of methods is called the context manager protocol.

🛠️ The Class Way: __enter__ and __exit__

Let’s build a timer. It starts counting when the block begins. It prints how long the block took when the block ends. First the full code, then we walk through it.

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()   # setup: remember the start time
        return self                         # this becomes the "as" value

    def __exit__(self, exc_type, exc_value, traceback):
        end = time.perf_counter()
        elapsed = end - self.start
        print(f"Block took {elapsed:.4f} seconds")
        # returning None (no return) means: do not hide any error

with Timer():
    total = sum(range(1_000_000))
    print(f"Sum is {total}")

Run it and you get something like this.

Sum is 499999500000
Block took 0.0123 seconds

Now the walkthrough, line by line.

• __enter__ runs first. It saves the start time on self, then returns self. The returned value is what as would catch if you wrote with Timer() as t:.
• The block runs your real work, the sum.
• __exit__ runs last, and it always runs. It measures the end time and prints how long it took.

Now look at the three odd arguments on __exit__. They are exc_type, exc_value, and traceback. These tell you whether the block ended because of an error.

• If the block finished normally, all three are None.
• If an error was raised inside the block, they hold the error type, the error object, and where it happened.

The key point is that the cleanup still runs when there is an error. Watch.

with Timer():
    print("starting work")
    raise ValueError("something broke")
    print("this line never runs")

We raise an error, but the timer still prints its line on the way out.

starting work
Block took 0.0001 seconds
Traceback (most recent call last):
  File "demo.py", line 18, in <module>
    raise ValueError("something broke")
ValueError: something broke

See how “Block took …” printed before the traceback? That is the promise. __exit__ always runs. The error then continues on its way, because we did not stop it.

🔁 A Real Resource: Always Closing

A timer is nice for learning. But the everyday use is opening a resource and making sure it closes. Imagine a simple connection object that must be closed when you finish.

class Connection:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        print(f"Opening connection to {self.name}")
        return self                       # hand the connection to the block

    def query(self, sql):
        print(f"Running: {sql}")

    def __exit__(self, exc_type, exc_value, traceback):
        print(f"Closing connection to {self.name}")

with Connection("users-db") as conn:
    conn.query("SELECT * FROM users")

The connection opens at the start. You use it inside. And it closes itself at the end.

Opening connection to users-db
Running: SELECT * FROM users
Closing connection to users-db

You never wrote a manual close() call, yet the connection closed. And it would still close even if query raised an error. That is the whole point. It is cleanup you cannot forget.

✨ The Easy Way: contextlib.contextmanager

Writing a whole class with two special methods feels like a lot for something small. Python gives you a shortcut. With the contextmanager decorator from the contextlib module, you write one normal function and use yield to split setup from cleanup.

Here is the timer again, the short way.

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.perf_counter()      # everything before yield = setup
    yield                            # the with block runs right here
    end = time.perf_counter()        # everything after yield = cleanup
    print(f"Block took {end - start:.4f} seconds")

with timer():
    total = sum(range(1_000_000))
    print(f"Sum is {total}")

The output is the same as the class version.

Sum is 499999500000
Block took 0.0119 seconds

Here is the way to read this.

• The code before yield is your __enter__.
• The yield itself is where your with block runs.
• The code after yield is your __exit__.

If you want an as value, you yield it. Here is the connection example written the short way, handing back a value.

from contextlib import contextmanager

@contextmanager
def connection(name):
    print(f"Opening connection to {name}")
    conn = {"name": name}            # pretend this is a real connection
    yield conn                       # this becomes the "as" value
    print(f"Closing connection to {name}")

with connection("users-db") as conn:
    print(f"Using {conn['name']}")

Opening connection to users-db
Using users-db
Closing connection to users-db

There is one catch with this short style. If an error happens in the block, it pops out at the yield line. So if your cleanup must run no matter what, wrap the yield in try/finally. The finally block always runs, error or not.

from contextlib import contextmanager

@contextmanager
def connection(name):
    print(f"Opening connection to {name}")
    try:
        yield {"name": name}         # block runs here
    finally:
        print(f"Closing connection to {name}")   # always runs

with connection("users-db") as conn:
    raise RuntimeError("query failed")

The connection still closes before the error continues.

Opening connection to users-db
Closing connection to users-db
Traceback (most recent call last):
  File "demo.py", line 13, in <module>
    raise RuntimeError("query failed")
RuntimeError: query failed

⚠️ Common Mistakes

A few traps catch people the first time they build their own.

• Forgetting to return something useful from __enter__. If you want with X() as y: to give you the object, you must return self. Forget it and y becomes None.

# ❌ Avoid: returns nothing, so "as f" is None
def __enter__(self):
    self.file = open("data.txt")

# ✅ Good: return the thing the block should use
def __enter__(self):
    self.file = open("data.txt")
    return self.file

• Putting cleanup after yield without try/finally. In the decorator style, an error in the block skips any plain line after yield. Your cleanup quietly does not run.

# ❌ Avoid: if the block errors, cleanup is skipped
@contextmanager
def session():
    s = open_session()
    yield s
    s.close()           # never runs if the block raised

# ✅ Good: finally guarantees cleanup
@contextmanager
def session():
    s = open_session()
    try:
        yield s
    finally:
        s.close()

• Returning True from __exit__ by habit. That hides every error raised in the block. Only return True when you truly mean “I handled this error, do not crash.”

✅ Best Practices

Small habits keep your context managers clean and predictable.

• Reach for the @contextmanager decorator for short, simple setup and cleanup. Use a full class when the object also needs other methods or holds a lot of state.
• Always guard the cleanup. Use try/finally in the decorator style. And remember a class __exit__ already runs on errors.
• Keep the setup and cleanup small and fast. A context manager is about managing a resource, not doing heavy work.
• Name the function or class after the resource it manages, like timer, connection, or temp_directory. Then the with line reads like plain English.

🧩 What You’ve Learned

• ✅ A context manager runs setup when a with block starts and cleanup when it ends. The cleanup runs even if an error is raised.
• ✅ The class way needs __enter__ (setup, returns the as value) and __exit__ (cleanup, gets error info).
• ✅ __exit__ returning True swallows the error. Returning None lets it continue.
• ✅ The @contextmanager decorator lets you write one function. Code before yield is setup and code after is cleanup.
• ✅ Wrap the yield in try/finally so cleanup runs even when the block errors.

🚀 What’s Next?

Now you can wrap setup and cleanup around your own objects. Next let’s look at a cleaner way to build classes that mostly just hold data. You get to skip writing all the boilerplate yourself.

Dataclasses