Create Custom Context Managers in Python

Creating your own context managers in Python is the step where you move from using the with statement for files and locks to designing your own abstractions for resource management. Every time you acquire a resource, use it, and release it, you are following a pattern that a context manager can formalize. The built-in context managers handle files, locks, and decimal precision. But your application has its own resources: database transactions, temporary directories, timer scopes, mock patches, log contexts, and network sessions. Learning to create custom context managers turns each of these ad-hoc try-finally patterns into a named, reusable abstraction that works with the with statement and guarantees cleanup regardless of exceptions.

The two previous articles, Python context managers explained and using the with statement in Python, established the protocol that context managers follow and how the with statement uses that protocol. This article shows you how to implement the protocol on your own classes. The approach is straightforward: define a class with an enter method that sets up the resource and an exit method that tears it down. The with statement calls these methods at the right times, and your class gets to define what setup and teardown mean for the specific resource it manages.

A class-based context manager has a natural lifecycle that mirrors how you think about resources. The init method can accept configuration parameters, the enter method acquires or prepares the resource before the block runs, the block itself uses the resource, and the exit method releases or cleans up the resource after the block ends, handling exceptions appropriately. This four-phase lifecycle (configure, acquire, use, release) maps cleanly onto Python's object model, and the resulting context manager is easy to test because each phase is a separate method with a well-defined responsibility.

A minimal custom context manager

The simplest possible context manager does nothing except prove that the protocol works. Its enter method prints a message and returns a value. Its exit method prints another message and returns None, which tells Python to propagate any exception that occurred. Here is the complete class:

python

class TraceContext:
    def __init__(self, name):
        self.name = name
 
    def __enter__(self):
        print(f"Entering {self.name}")
        return self
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        print(f"Exiting {self.name}")
        return False

The init method stores a name for identification. The enter method announces that the block is beginning and returns the instance itself, which becomes available through the as clause. The exit method announces that the block is ending and returns False, which means "do not suppress exceptions." The three exception parameters (type, value, and traceback) are all None when the block completes normally, and they carry the exception information when an exception occurred. Returning False or None tells Python to propagate the exception; returning True tells Python to suppress it.

Using this context manager is identical to using any built-in one. The with statement creates an instance, calls enter before the block, runs the block, and calls exit afterwards. Any code that can use a built-in context manager can use this custom one without modification:

python

with TraceContext("database-query"):
    print("Running query")

The output shows the enter message, then the block's message, then the exit message. If the block raised an exception, the exit message would still appear because exit runs in the finally clause that the with statement constructs. The exception would then propagate because exit returned False.

A practical example: managed temporary directory

A more useful custom context manager manages a temporary directory that is created before the block and deleted afterwards. This pattern is common in testing, where tests need a clean working directory that is automatically cleaned up regardless of whether the test passes or fails. The standard library's tempfile module provides TemporaryDirectory for this purpose, and building a simplified version illustrates the pattern.

The context manager creates a unique temporary directory in its enter method, returns the path so the block can use it, and recursively deletes the directory in its exit method. The cleanup must happen even if the block raises an exception, and the exit method must handle the case where the directory was already deleted or never created:

python

import os
import tempfile
import shutil
 
class TempDir:
    def __init__(self, prefix="tmp"):
        self.prefix = prefix
        self.path = None
 
    def __enter__(self):
        self.path = tempfile.mkdtemp(prefix=self.prefix)
        return self.path
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.path and os.path.exists(self.path):
            shutil.rmtree(self.path)
        return False

The enter method creates the directory using mkdtemp from the tempfile module, which generates a unique name and creates the directory atomically. The path is stored as an instance attribute so exit can access it, and it is also returned through the as clause so the block can use it. The exit method checks that the path was set and that the directory still exists before attempting to delete it. The rmtree function from shutil recursively removes the directory and all its contents.

Using this context manager, a test can create files in a temporary directory, verify their contents, and know that the directory will be cleaned up afterwards without writing any cleanup code in the test body:

python

with TempDir(prefix="test-") as tmp:
    filepath = os.path.join(tmp, "data.txt")
    with open(filepath, "w") as f:
        f.write("test data")
    assert os.path.exists(filepath)

When the block ends, the temporary directory and the file inside it are deleted. If the assertion fails, the directory is still deleted because the with statement calls exit even when the block raises an exception. The test author does not need to write a try-finally block or remember to call rmtree. The context manager encapsulates the cleanup contract.

Handling exceptions in exit

The exit method receives three arguments that describe any exception that occurred in the block. The first argument is the exception type (a class like ValueError or ZeroDivisionError). The second is the exception instance. The third is the traceback object. When no exception occurred, all three are None. This information lets the exit method decide how to handle the situation: it can log the exception and propagate it, suppress specific exceptions, or perform different cleanup depending on whether the block succeeded or failed.

A context manager that opens a database transaction might use the exception information to decide whether to commit or roll back. If exit receives None for all three exception arguments, the transaction committed successfully and the manager can call commit. If exit receives an exception type, something went wrong and the manager should call rollback. Here is the pattern:

python

class Transaction:
    def __init__(self, connection):
        self.conn = connection
 
    def __enter__(self):
        self.conn.begin()
        return self.conn
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            self.conn.commit()
        else:
            self.conn.rollback()
        return False

The enter method begins the transaction. The block runs with the connection and performs database operations. The exit method checks whether an exception occurred by testing whether exc_type is None. If no exception occurred, the transaction is committed. If an exception occurred, the transaction is rolled back, and the exception propagates because exit returns False. The caller does not need to write commit or rollback logic; the context manager handles both paths correctly.

This pattern of branching on the presence of an exception in exit is the foundation of resource managers that need to distinguish between successful and failed operations. File copy operations, network transfers, batch processing jobs, and any multi-step operation that must be atomic can use this pattern to guarantee that partial work is cleaned up on failure.

Returning True to suppress exceptions

The exit method can suppress an exception by returning True. When exit returns a truthy value, Python does not propagate the exception. Execution continues with the statement after the with block. This capability is used sparingly because suppressing exceptions can hide bugs, but there are legitimate cases. The contextlib.suppress context manager in the standard library is a built-in example: it catches and suppresses specified exception types.

A context manager that retries a failing operation a fixed number of times might suppress the intermediate failures and only propagate the final one. Or a context manager that logs and ignores non-critical errors might suppress exceptions that represent recoverable conditions. The key is that the decision to suppress should be explicit and documented, and the suppressed exception should be logged or otherwise recorded so it is not silently lost.

The next article covers Python context managers with contextlib, the standard library module that provides utilities for working with context managers. It includes the contextmanager decorator, which lets you write a context manager as a generator function instead of a class, and several pre-built context managers for common patterns like suppressing exceptions and redirecting standard output.