Skip to Content
Course content

96: Writing Your Own Context Manager

Click on the "Edit" button in the top corner of the screen to edit your slide content.

By now, you've used the with statement plenty of times—probably mostly for opening files. It's one of those Python features that just feels "right" because it guarantees that resources get cleaned up, even if your code crashes halfway through. But eventually, you'll hit a scenario where Python's built-in managers aren't enough. You'll find yourself needing to manage something specific to your own business logic, like a database transaction, a temporary configuration change, or a network socket.

The repetitive try-finally dance

Let's say we have a global configuration object that controls how our app logs data. Occasionally, for a specific block of code, we want to force the log level to DEBUG regardless of the global setting, but we must ensure it flips back to the original state afterward. The naive way to handle this is a try...finally block:

import config

def process_data():
    old_level = config.LOG_LEVEL
    config.LOG_LEVEL = "DEBUG"
    try:
        # Do some complex work here
        print("Processing with debug logs...")
        # Imagine an exception happens here
    finally:
        config.LOG_LEVEL = old_level

This works, but it's brittle. If you need to do this in ten different places across your codebase, you're repeating the same setup and teardown logic over and over. If you forget the finally block just once, and an exception hits, your entire application stays in DEBUG mode for the rest of its lifecycle. That's a nightmare to debug in production. I've spent far too many hours chasing bugs caused by "leaked" states exactly like this.

Encapsulating the lifecycle

The professional way to solve this is to wrap that logic into a context manager. A context manager is just an object that defines two magic methods: __enter__ and __exit__. The __enter__ method handles the setup, and __exit__ handles the cleanup.

class TempLogLevel:
    def __init__(self, new_level):
        self.new_level = new_level
        self.old_level = None

    def __enter__(self):
        self.old_level = config.LOG_LEVEL
        config.LOG_LEVEL = self.new_level
        return self # This is what 'as' assigns to

    def __exit__(self, exc_type, exc_val, exc_tb):
        config.LOG_LEVEL = self.old_level
        # Returning False (default) lets any exceptions propagate up

Now, your business logic becomes clean and declarative. You don't have to remember the teardown; the language handles it for you:

with TempLogLevel("DEBUG"):
    process_data()
# The log level is guaranteed to be restored here

The beauty here is that __exit__ receives the exception details if something went wrong. If you wanted your manager to "swallow" a specific error and keep the program running, you could return True from __exit__. Normally, though, you'll just let the exception bubble up after cleaning up your state.

When you don't want to write a whole class

Writing a full class for a simple setup/teardown is sometimes overkill. Python provides a shortcut in the contextlib module called the @contextmanager decorator. This allows you to turn a generator function into a context manager. It's much more concise and, in my opinion, easier to read for simple cases.

from contextlib import contextmanager

@contextmanager
def temp_log_level(new_level):
    old_level = config.LOG_LEVEL
    config.LOG_LEVEL = new_level
    try:
        yield # This is where the code inside the 'with' block executes
    finally:
        config.LOG_LEVEL = old_level

Everything before the yield is your __enter__; everything after is your __exit__. The try...finally is still necessary inside the generator to ensure the cleanup happens even if the yielded block raises an error. Between the class-based approach and the decorator, the decorator is almost always the right choice unless you need to maintain complex internal state within the manager object itself.




📋 Practical Task

Build a "Execution Timer" Context Manager

In professional profiling, we often need to know exactly how long a specific block of code takes to run without littering our functions with time.time() calls. Your task is to create a context manager that measures the elapsed time of the code block it wraps.

Requirements:

  • Create a context manager named ExecutionTimer (you may use either the class-based approach or the @contextmanager decorator).
  • The manager should accept a string label as an argument (e.g., "Database Query").
  • When the block exits, it should print the label and the total time elapsed in seconds.
  • Ensure that the time is printed even if the code inside the with block raises an exception.
# Expected Usage:
with ExecutionTimer("Heavy Computation"):
    # simulate some work
    sum(i*i for i in range(1000000))

# Output: Heavy Computation took 0.123 seconds
Rating
0 0

There are no comments for now.

to be the first to leave a comment.