Custom Context Manager (Class-Based)
A class-based context manager defines `__enter__` and `__exit__` so the object can be used in a `with` block. It is the right shape when the resource has setup, teardown, AND state you want to expose to the body (file handles, DB connections, locks). This entry shows the basic skeleton, an exception-aware variant, and the gotchas around return values from `__exit__`.
217 views
5
class Timer:
"""Measure how long the with-block took, in seconds."""
def __enter__(self):
import time
self.start = time.perf_counter()
return self # value bound by `as` in the with-statement
def __exit__(self, exc_type, exc_value, traceback):
import time
self.elapsed = time.perf_counter() - self.start
# Return None / False so any exception is re-raised.
with Timer() as t:
total = sum(i * i for i in range(100_000))
print('sum:', total)
print(f'elapsed: {t.elapsed:.4f}s')__enter__ runs at the top of the with block; whatever it returns is bound to the as target. __exit__ always runs when the block ends, whether by falling off the bottom or by an exception. The three arguments (exc_type, exc_value, traceback) are all None on a clean exit and describe the exception otherwise. Returning None (or any falsy value) lets the exception propagate, which is what 99% of context managers want. The pattern is the right tool whenever 'do X before, undo X after' has to survive every exit path.
class IgnoreNotFound:
"""Run a block and quietly swallow FileNotFoundError."""
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
# Returning True tells Python: 'I handled it, do not re-raise'.
if exc_type is FileNotFoundError:
print(f'(ignored: {exc_value})')
return True
return False # everything else propagates
with IgnoreNotFound():
open('this-file-definitely-does-not-exist.txt')
print('made it past the with-block')
try:
with IgnoreNotFound():
raise ValueError('different exception')
except ValueError as exc:
print('ValueError still propagates:', exc)__exit__ is the only place a context manager can decide whether to swallow an in-flight exception. Returning a truthy value tells Python to treat the exception as handled, which is how contextlib.suppress is implemented. Use this power sparingly: silencing exceptions hides bugs, so always check exc_type and only swallow the one you mean to. Returning False (or None) for every other exception keeps the default 'propagate' behaviour for unexpected errors.
class TrackedConnection:
"""Pretend connection that must be closed even on error."""
open_count = 0
def __init__(self, name):
self.name = name
self.closed = False
def __enter__(self):
TrackedConnection.open_count += 1
print(f'open {self.name} (count={TrackedConnection.open_count})')
return self
def __exit__(self, exc_type, exc_value, traceback):
self.closed = True
TrackedConnection.open_count -= 1
print(f'close {self.name} (count={TrackedConnection.open_count})')
# Do not swallow the exception; let it propagate.
return False
def query(self, sql):
if self.closed:
raise RuntimeError('connection already closed')
return f'{self.name} ran: {sql}'
# Happy path.
with TrackedConnection('primary') as c:
print(c.query('SELECT 1'))
# Sad path: the connection still gets closed.
try:
with TrackedConnection('replica') as c:
raise RuntimeError('boom')
except RuntimeError as exc:
print('caught:', exc)
print('open_count after both blocks:', TrackedConnection.open_count)The classic use case is anything with paired open / close semantics: sockets, files, DB cursors, mutexes. The class form shines when the resource has methods you want to call inside the block (c.query(...)) and state worth tracking on the instance (c.closed). Even when the body raises, __exit__ still runs first, so cleanup is guaranteed without a hand-written try / finally. The open_count invariant going back to zero in both branches is the proof: the with block survives every exit path.