Python Custom Context Manager: How to Implement

A custom context manager in Python lets you control what happens before and after a block of code runs. You might need one when the built-in context managers don't cover your resource lifecycle, or when you want to encapsulate setup and teardown logic that repeats across your codebase. This article explains how to implement a python custom context manager using the two standard approaches: the enter/exit protocol and the contextlib.contextmanager decorator.

The Two Ways to Define a Custom Context Manager

Python provides two primary mechanisms for creating custom context managers: implementing the enter and exit methods on a class, or using the contextlib.contextmanager decorator on a generator function. Both integrate with the with statement, but they differ in structure and use cases.

Implementing with enter and exit

The class-based approach requires you to define a class with enter and exit methods. The enter method is called when the with block is entered, and its return value is bound to the target variable after as. The exit method is called when the block exits, regardless of whether an exception occurred.

class ManagedFile:
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode

    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()
        return False

In this example, enter opens the file and returns it, so you can use it inside the with block. exit closes the file. The three arguments to exit carry exception information: exc_type is the exception class, exc_val is the instance, and exc_tb is the traceback. If no exception occurred, all three are None. Returning False (the default) means the exception, if any, propagates after exit finishes. Returning True suppresses the exception, which is rarely what you want.

Using contextlib.contextmanager for Simpler Cases

The contextlib.contextmanager decorator lets you write a context manager as a generator function with a single yield. Code before the yield runs on entry; code after the yield runs on exit. This is often more readable for simple setup/teardown logic.

from contextlib import contextmanager

@contextmanager
def managed_file(filename, mode):
    file = open(filename, mode)
    try:
        yield file
    finally:
        file.close()

The try/finally ensures cleanup even if an exception occurs inside the with block. The value yielded is what gets bound to the as target. If you need to handle exceptions, you can wrap the yield in a try/except block and inspect the exception type.

Handling Exceptions and Propagation

Both approaches allow you to control exception propagation. In the class-based version, exit can return True to swallow an exception, but this is usually a code smell. In the generator-based version, if an exception occurs inside the with block, it is thrown into the generator at the yield point. You can catch it and decide whether to suppress it or re-raise.

@contextmanager
def suppress_specific_error():
    try:
        yield
    except ValueError:
        print("ValueError suppressed")

Here, a ValueError raised inside the with block is caught and not propagated. Other exceptions propagate normally. This pattern is useful for selectively ignoring certain failures, but you should be explicit about which exceptions you handle.

Practical Use Cases for Custom Context Managers

Custom context managers shine when you need to enforce a consistent setup and teardown sequence. Common examples include timing code blocks, temporarily changing environment variables, managing database transactions, or switching the current working directory.

import time
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"Elapsed: {elapsed:.4f} seconds")

Using this context manager, you can measure any block of code without duplicating timing logic. Similarly, a context manager that changes the working directory and restores it afterward prevents state leakage between tests or functions.

Performance and Overhead Considerations

The overhead of a context manager is minimal. Entering and exiting a with block involves two method calls (or a generator yield and resume). For most applications, this is negligible. However, if you are writing a context manager that wraps a very tight loop, the extra function calls can add up. In such cases, consider whether the readability and safety gains outweigh the microsecond-level cost. The generator-based approach has slightly more overhead than the class-based one because it involves a generator frame, but again, this is rarely a bottleneck.

Common Mistakes and Pitfalls

One frequent mistake is forgetting to return a value from enter. If you don't return anything, the as target gets None, which can lead to confusing AttributeError later. Another issue is not using try/finally in the generator-based version, which means cleanup code may not run if an exception occurs. Also, be careful not to yield more than once in a contextmanager-decorated function; doing so raises a RuntimeError. Finally, remember that exit should not re-raise exceptions unless you have a specific reason; let the Python runtime handle propagation.

When to Choose Which Approach

Use the class-based approach when you need to maintain state across multiple context entries, or when the context manager is complex enough to benefit from methods and attributes. Use contextlib.contextmanager when the logic is linear and fits naturally in a generator function. The decorator is often shorter and easier to read for simple cases. If you are building a library that exposes a context manager, the class-based approach may be more explicit and easier to document.

Advanced: Context Managers for Asynchronous Code

Python's async with statement requires an asynchronous context manager, which you implement with aenter and aexit methods, or by decorating an async generator with @asynccontextmanager. The same principles apply, but the methods are coroutines. This is useful when acquiring and releasing resources asynchronously, such as database connections or network sessions.