How to write custom context managers in Python
Writing a custom context manager is straightforward: create a class with __enter__ and __exit__ methods, implement acquisition logic in __enter__, and cleanup logic in __exit__. Custom context managers are powerful because they encapsulate resource management, making your code more readable, maintainable, and exception-safe. This article walks through the design patterns, common mistakes, and real-world examples.
The best context managers are single-purpose: one responsibility, one resource type, one clear cleanup contract. By following this principle and understanding the protocol deeply, you can write context managers that are not only correct but also intuitive and reusable across projects.
Basic Custom Context Manager Template
Every custom context manager follows this template: define __enter__ to acquire, and __exit__ to cleanup. The minimal implementation requires only these two methods.
# Template: minimal custom context manager
class MyContextManager:
def __init__(self, resource_config):
"""Initialize configuration, but do not acquire resources yet."""
self.config = resource_config
self.resource = None
def __enter__(self):
"""Acquire the resource. Called on entering 'with' block."""
print(f"Acquiring resource with config: {self.config}")
# Simulate acquiring a resource
self.resource = f"Resource({self.config})"
return self.resource # Return what will be bound to 'as' variable
def __exit__(self, exc_type, exc_val, exc_tb):
"""Clean up the resource. Called on exiting 'with' block."""
print(f"Releasing resource: {self.resource}")
self.resource = None
# Return False to let exceptions propagate
return False
# Using the custom context manager
with MyContextManager("my_config") as resource:
print(f"Using: {resource}")
# Output:
# Acquiring resource with config: my_config
# Using: Resource(my_config)
# Releasing resource: Resource(my_config)
Design Principle: Acquire in enter, Release in exit
The clearest context managers follow a simple rule: never acquire resources in __init__. Instead, acquire them in __enter__. This separation ensures that the context manager is cheap to create, and resources are only acquired when the with statement executes. If __enter__ is never called, no resources are wasted.
# Example 1: Resource acquisition order
class DatabaseConnection:
def __init__(self, connection_string):
# Store configuration, but do not connect yet
self.connection_string = connection_string
self.connection = None
print(f"Context manager created (no connection yet)")
def __enter__(self):
# Acquire the resource
print(f"Connecting to {self.connection_string}")
self.connection = f"DB connection: {self.connection_string}"
return self.connection
def __exit__(self, exc_type, exc_val, exc_tb):
# Release the resource
if self.connection:
print(f"Closing {self.connection}")
self.connection = None
return False
# Create the context manager
db_cm = DatabaseConnection("postgres://localhost/mydb")
print("Context manager created")
# Enter the with block to actually connect
with db_cm as conn:
print(f"Connected: {conn}")
# Output:
# Context manager created (no connection yet)
# Context manager created
# Connecting to postgres://localhost/mydb
# Connected: DB connection: postgres://localhost/mydb
# Closing DB connection: postgres://localhost/mydb
Handling Partial Acquisition Failures
If resource acquisition fails partway through, __exit__ is never called. So you must decide: either keep __enter__ atomic (all-or-nothing), or handle cleanup for partial acquisitions inside __enter__ itself.
# Example 2: Atomic acquisition vs. partial cleanup
class MultiResourceManager:
def __init__(self):
self.resource_a = None
self.resource_b = None
self.resource_c = None
def __enter__(self):
try:
# Acquire resource A
self.resource_a = "Resource A"
print(f"Acquired: {self.resource_a}")
# Acquire resource B
self.resource_b = "Resource B"
print(f"Acquired: {self.resource_b}")
# Attempt to acquire resource C (might fail)
if False: # Simulate failure
raise RuntimeError("C unavailable")
self.resource_c = "Resource C"
print(f"Acquired: {self.resource_c}")
return self
except Exception as e:
# If acquisition fails partway, clean up what was acquired
print(f"Acquisition failed: {e}. Cleaning up partial resources.")
if self.resource_a:
print(f"Releasing {self.resource_a}")
self.resource_a = None
if self.resource_b:
print(f"Releasing {self.resource_b}")
self.resource_b = None
# Re-raise to signal failure
raise
def __exit__(self, exc_type, exc_val, exc_tb):
print("Cleanup called")
for attr in ['resource_a', 'resource_b', 'resource_c']:
if getattr(self, attr):
print(f"Releasing {getattr(self, attr)}")
setattr(self, attr, None)
return False
# Test successful acquisition
print("--- Successful acquisition ---")
with MultiResourceManager() as m:
print(f"Working with {m.resource_a}, {m.resource_b}, {m.resource_c}")
# Output:
# --- Successful acquisition ---
# Acquired: Resource A
# Acquired: Resource B
# Acquired: Resource C
# Working with Resource A, Resource B, Resource C
# Cleanup called
# Releasing Resource A
# Releasing Resource B
# Releasing Resource C
Returning Different Objects from enter
__enter__ can return anything: self, a wrapper object, configuration, or even None. Returning self is common when the context manager is the resource. Returning a different object is useful when you want to expose a limited API or hide implementation details.
# Example 3: Returning a wrapper instead of self
class FileWriterManager:
def __init__(self, filename):
self.filename = filename
self.file = None
self.wrapper = None
def __enter__(self):
self.file = open(self.filename, 'w')
# Return a wrapper with a limited API
self.wrapper = FileWriterWrapper(self.file)
return self.wrapper
def __exit__(self, exc_type, exc_val, exc_tb):
if self.file:
self.file.close()
return False
class FileWriterWrapper:
"""A limited API that only allows writing, not seeking or other operations."""
def __init__(self, file):
self.file = file
def write(self, text):
self.file.write(text)
import tempfile
import os
with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp:
tmp_name = tmp.name
try:
with FileWriterManager(tmp_name) as writer:
writer.write("Hello from wrapper")
# writer.seek(0) # Would fail: FileWriterWrapper has no seek
with open(tmp_name, 'r') as f:
print(f"File contents: {f.read()}")
finally:
os.unlink(tmp_name)
# Output:
# File contents: Hello from wrapper
Optional Cleanup and Stateful exit
Context managers can track state during acquisition and use that state during cleanup. This enables sophisticated cleanup strategies, such as rolling back if the block raised an exception.
# Example 4: Stateful cleanup for transactions
class TransactionManager:
def __init__(self, db_name):
self.db_name = db_name
self.transaction_id = None
self.committed = False
def __enter__(self):
# Simulate starting a transaction
import time
self.transaction_id = f"txn_{int(time.time() * 1000)}"
print(f"[{self.db_name}] Started transaction: {self.transaction_id}")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_type is not None:
# An exception occurred; rollback
print(f"[{self.db_name}] Rolling back {self.transaction_id} due to {exc_type.__name__}")
self.committed = False
else:
# No exception; commit
print(f"[{self.db_name}] Committing {self.transaction_id}")
self.committed = True
self.transaction_id = None
return False
# Test successful transaction
print("--- Successful commit ---")
with TransactionManager("mydb") as txn:
print(f"Transaction {txn.transaction_id} active")
print()
# Test rollback on exception
print("--- Rollback on exception ---")
try:
with TransactionManager("mydb") as txn:
print(f"Transaction {txn.transaction_id} active")
raise ValueError("Operation failed")
except ValueError:
print("Exception handled by caller")
# Output:
# --- Successful commit ---
# [mydb] Started transaction: txn_1717327234567
# Transaction txn_1717327234567 active
# [mydb] Committing txn_1717327234567
#
# --- Rollback on exception ---
# [mydb] Started transaction: txn_1717327234568
# Transaction txn_1717327234568 active
# [mydb] Rolling back txn_1717327234568 due to ValueError
# Exception handled by caller
Defensive Cleanup: Always Success
The most robust context managers implement defensive cleanup: they ensure __exit__ succeeds regardless of what happened, and they log errors rather than propagating them during cleanup.
# Example 5: Defensive cleanup pattern
class RobustContextManager:
def __init__(self):
self.resources = []
def __enter__(self):
self.resources = ["A", "B", "C"]
print(f"Acquired: {self.resources}")
return self
def __exit__(self, exc_type, exc_val, exc_tb):
# Release all resources, even if cleanup fails
for resource in reversed(self.resources):
try:
print(f"Releasing {resource}")
# Simulate occasional cleanup failures
if resource == "B" and exc_type is None:
raise RuntimeError(f"Failed to release {resource}")
self.resources.remove(resource)
except Exception as e:
# Log but don't propagate cleanup errors
print(f"Warning: cleanup failed for {resource}: {e}")
if exc_type is None:
# If no original exception, propagate this cleanup error
raise
return False
# Test normal cleanup
print("--- Normal cleanup ---")
with RobustContextManager() as m:
print("Working...")
print()
# Test cleanup with original exception
print("--- Cleanup with original exception ---")
try:
with RobustContextManager() as m:
raise ValueError("Original error")
except ValueError as e:
print(f"Caught: {e}")
# Output:
# --- Normal cleanup ---
# Acquired: ['A', 'B', 'C']
# Working...
# Releasing C
# Releasing B
# Warning: cleanup failed for B: Failed to release B
# Releasing A
#
# --- Cleanup with original exception ---
# Acquired: ['A', 'B', 'C']
# Releasing C
# Releasing B
# Releasing A
# Caught: ValueError: Original error
Key Takeaways
- Design context managers with a single responsibility: acquire one resource type in
__enter__and release it in__exit__. - Never acquire resources in
__init__; only store configuration. Acquire in__enter__. - If acquisition fails partway, clean up partial resources inside
__enter__itself before re-raising. - Return
selffrom__enter__when the context manager is the resource; return a wrapper when you want to hide implementation. - Implement defensive cleanup in
__exit__: catch and log errors without suppressing the original exception.
Frequently Asked Questions
Should I acquire resources in init or enter?
Always acquire in __enter__. Acquiring in __init__ wastes resources if the context manager is created but the with block is never entered. It also makes testing harder because you must manually clean up in tests.
Can I use inheritance with context managers?
Yes. You can inherit __enter__ and override __exit__, or vice versa. However, if you override either method, call super().__enter__() and super().__exit__() to preserve cleanup contracts. Use ExitStack (see next article) for complex cleanup hierarchies.
What if I need to acquire multiple independent resources?
Either create separate context managers and nest them, or create a single context manager that acquires all and releases all in reverse order (LIFO). For complex scenarios, use contextlib.ExitStack.
How do I test a custom context manager?
Test the context manager directly: call __enter__, verify the return value and side effects, then call __exit__ with and without exceptions. Use pytest fixtures or helper functions to assert cleanup happened.
Can a context manager be reused (used in multiple with blocks)?
No, not by default. Each with statement creates a new instantiation or reuses an existing context manager and calls __enter__ / __exit__ once. If you need reusable resources, wrap the context manager instantiation in a factory function or use a connection pool with a context manager API.