You already know that the with statement is the gold standard in Python for managing resources. Whether you're closing a file or releasing a lock, it ensures that cleanup happens regardless of whether the code inside the block succeeds or crashes. But for a long time, the only way to create your own context manager was to define a class with __enter__ and __exit__ methods. Honestly? It's a slog. It feels like you're writing more infrastructure than actual logic.

The boilerplate tax of class-based managers

Let's say you're writing a test suite and you need to temporarily change an environment variable—maybe an API key or a debug flag—and then restore it to its original state immediately after the test finishes. If you do this the "official" class-based way, it looks like this:

import os

class TempEnvVar:
    def __init__(self, key, value):
        self.key = key
        self.value = value
        self.old_value = None

    def __enter__(self):
        self.old_value = os.environ.get(self.key)
        os.environ[self.key] = self.value
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.old_value is None:
            os.environ.pop(self.key, None)
        else:
            os.environ[self.key] = self.old_value

Now, look at how much of that is just " ceremony." You have to manage the state in __init__, remember to return self (or something else) in __enter__, and deal with those three clunky exception arguments in __exit__, even if you don't plan on using them. It spreads the setup and teardown logic across different methods, which makes the flow harder to follow at a glance.

Flattening the logic with a generator

This is where contextlib.contextmanager comes in. It's a decorator that lets you turn a simple generator function into a full-blown context manager. Instead of splitting your logic into two methods, you write it as a single linear sequence. Everything before the yield is your setup; everything after the yield is your teardown.

import os
from contextlib import contextmanager

@contextmanager
def temp_env_var(key, value):
    old_value = os.environ.get(key)
    os.environ[key] = value
    try:
        yield
    finally:
        if old_value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = old_value

I love this approach because it reads like a story. "Save the old value, set the new one, let the user do their thing, and then put it back." You've eliminated the class overhead and the __exit__ signature. When you use with temp_env_var("API_KEY", "test_val"):, Python runs the function up until the yield`, then executes the block inside your with statement, and finally resumes the function after the yield.

The danger of skipping the finally block

There is one massive "gotcha" here that I see developers miss all the time. In a class-based manager, __exit__ is guaranteed to run. In a generator-based manager, if an exception occurs inside the with block, that exception is "thrown" back into your generator at the point of the yield.

If you don't wrap your yield in a try...finally block, and an error occurs in the with block, the code after your yield will never execute. In our environment variable example, if the test fails, your environment variable stays changed, potentially breaking every other test in your suite. Using finally ensures that the cleanup happens regardless of whether the protected code succeeded or raised an exception. It's not optional; it's the only way to make the decorator safe.

📋 Practical Task

Build a Precision Block Timer

You need to profile a specific section of a legacy codebase to see exactly how long a heavy computation is taking. Instead of manually calling time.perf_counter() before and after the block, create a context manager using @contextlib.contextmanager that handles the timing for you.

Requirements:

• Create a function called block_timer decorated with @contextmanager.
• The manager should start a timer when entering the block.
• It should yield a value (you can yield None or the start time).
• It must use a try...finally block to ensure that the end time is captured and the total duration is printed to the console, even if the code inside the with block raises an exception.
• The output should be in the format: "Block execution time: [X] seconds".

Test your implementation with this snippet:

# This should print the time taken and still show the error
try:
    with block_timer():
        import time
        time.sleep(0.5)
        raise RuntimeError("Something went wrong!")
except RuntimeError:
    pass