Creating Context Managers with `@contextlib.contextmanager`

1. Two Paths to a Context Manager: Class vs. Function

To appreciate the elegance of the decorator-based approach, it's useful to first see the more explicit, class-based method. This provides the "bad vs. good" comparison you appreciate and clarifies what's happening under the hood.

A class becomes a context manager by implementing two special "dunder" methods:

• __enter__(self): The setup logic. Runs when the with block is entered. Its return value is bound to the as variable.
• __exit__(self, exc_type, exc_value, exc_tb): The teardown logic. Runs when the with block is exited, guaranteed. It receives exception details if one occurred.

As you can see, the class-based approach is perfectly functional but involves a fair amount of ceremony. For most custom context managers, there is a much cleaner way.

2. The Pythonic Way: @contextlib.contextmanager

The contextlib module is part of Python's "batteries-included" standard library, and it contains the powerful @contextmanager decorator. This decorator transforms a simple generator function into a full-fledged context manager.

The structure is straightforward and elegant:

1. Import contextmanager from contextlib.
2. Define a generator function and decorate it with @contextmanager.
3. The code before the yield statement is your setup logic (the __enter__ part).
4. The yield statement passes control back to the with block. The value you yield becomes available to the as variable.
5. The code after the yield statement is your teardown logic (the __exit__ part).

Let's see this in action.

The core pattern you just saw is essential:

from contextlib import contextmanager

@contextmanager
def my_context_manager():



    # 1. Setup logic runs here
    print("Setting up...")
    
    try:



        # 2. Yield to the 'with' block
        yield
    finally:



        # 3. Teardown logic runs here, guaranteed
        print("Tearing down...")




# Usage
with my_context_manager():
    print("Inside the 'with' block.")

This is significantly more readable than the class-based version for most use cases. The logical flow—setup, do the work, teardown—is linear and easy to follow.

3. A Practical Example: Temporarily Changing Directory

Re-implementing open() is a good learning exercise, but let's build something genuinely useful that isn't already a built-in. A common task in scripts is to temporarily change to a different directory, perform some actions, and then reliably change back, even if errors occur. This is a perfect candidate for a context manager.

4. Yielding a Value to the as Variable

In the change_dir example, we didn't need to pass an object into the with block. But what if you do? For example, you might want your context manager to provide a database connection, a temporary object, or a logging handle.

To do this, you simply pass a value with the yield statement.

5. When to Stick with a Class

While @contextmanager is the preferred tool for its simplicity, it's not a silver bullet. For highly complex context managers, especially those that need to manage significant internal state or expose methods on the context manager object itself, a class can sometimes be clearer.

A general rule of thumb:

• If your __enter__ and __exit__ logic is straightforward, use @contextmanager.
• If your context manager needs to maintain complex state across the with block, or if you find yourself returning self from __enter__ to call methods on it, a class might be a better fit.

Exercise: Create a Simple Timer Context Manager

Now it's your turn to build a context manager from scratch. Your task is to create a function time_block that can be used to time the execution of any block of code.

Requirements:

1. The function should be a context manager created with the @contextmanager decorator.
2. Upon entering the with block, it should record the start time using time.perf_counter().
3. Upon exiting the with block, it should record the end time, calculate the duration, and print a formatted message like Block executed in: 0.1234 seconds.
4. The cleanup (printing the time) must happen even if an exception occurs inside the with block.

Starter Code:

import time
from contextlib import contextmanager




# Your implementation here
@contextmanager
def time_block():



    # TODO: Implement the setup, yield, and teardown logic
    pass




# --- Example Usage ---

print("--- Timing a successful block ---")
with time_block():



    # Simulate some work
    time.sleep(0.5)

print("\n--- Timing a block with an error ---")
try:
    with time_block():
        time.sleep(0.2)
        raise ValueError("Something went wrong!")
except ValueError as e:
    print(f"Caught expected error: {e}")




# Check that the timer message still printed for the failed block.

import time
from contextlib import contextmanager

@contextmanager
def time_block():
    """A context manager to time a block of code."""
    start_time = time.perf_counter()
    try:



        # Yield control back to the with block
        yield
    finally:



        # This code is guaranteed to run
        end_time = time.perf_counter()
        elapsed_time = end_time - start_time
        print(f"Block executed in: {elapsed_time:.4f} seconds")




# --- Example Usage ---

print("--- Timing a successful block ---")
with time_block():



    # Simulate some work
    time.sleep(0.5)

print("\n--- Timing a block with an error ---")
try:
    with time_block():
        time.sleep(0.2)
        raise ValueError("Something went wrong!")
except ValueError as e:
    print(f"Caught expected error: {e}")

Conclusion

You have now learned how to create your own context managers, a powerful feature for building robust and reusable APIs. By wrapping setup and teardown logic in a generator and using the @contextlib.contextmanager decorator, you can reduce boilerplate and make your code's intent clearer.

Key Takeaways:

• Custom context managers automate setup and teardown logic for your own objects.
• The @contextlib.contextmanager decorator turns a generator function into a context manager.
• The code before yield is for setup, and the code after yield is for teardown.
• Wrapping the yield in a try...finally block is the standard pattern to guarantee cleanup.
• The value passed to yield is what gets assigned to the as variable in a with statement.
• For simple, stateless context managers, the decorator is more Pythonic and concise than a full class.

In our next lesson, we will continue exploring the contextlib module by looking at other useful utilities like contextlib.suppress to cleanly ignore specific exceptions and contextlib.ExitStack to manage a dynamic number of contexts.