Custom Context Manager via contextlib
`@contextlib.contextmanager` turns a generator into a context manager: code before the `yield` is the setup, the yielded value is bound by `as`, and code after the `yield` is the teardown. It removes most of the class boilerplate when you do not need shared state. This entry covers the basic pattern, exception handling with try/finally, and `contextlib.suppress` as a one-liner.
968 views
17
from contextlib import contextmanager
@contextmanager
def temporary_setting(settings, key, value):
"""Set settings[key] for the duration of the with-block, then restore."""
original = settings.get(key)
had_key = key in settings
settings[key] = value
try:
yield settings # whatever you yield is bound by `as`
finally:
if had_key:
settings[key] = original
else:
settings.pop(key, None)
config = {'log_level': 'INFO'}
with temporary_setting(config, 'log_level', 'DEBUG') as cfg:
print('inside:', cfg['log_level']) # DEBUG
print('after :', config['log_level']) # INFO (restored)
with temporary_setting(config, 'feature_x', True):
print('inside:', config.get('feature_x')) # True
print('after :', config.get('feature_x')) # None (key removed)The generator runs once: code before yield is the __enter__ body, the yielded object is bound by as, and code after yield is the __exit__ body. Wrapping the post-yield code in try / finally is mandatory when teardown must run on exceptions too: without it, an exception from the body skips the cleanup. The decorator handles the boilerplate of building a class with __enter__ and __exit__, so the implementation reads top-to-bottom like normal code. Use this pattern for stateless, run-once setup-and-teardown helpers.
from contextlib import contextmanager
@contextmanager
def section(title):
print(f'==> {title}: start')
try:
yield
except Exception as exc:
print(f'==> {title}: failed ({type(exc).__name__}: {exc})')
raise # re-raise so callers still see the error
else:
print(f'==> {title}: ok')
finally:
print(f'==> {title}: cleanup')
with section('happy path'):
print(' ... doing work')
print()
try:
with section('sad path'):
raise ValueError('oops')
except ValueError:
print('caller saw the ValueError')The try / except / else / finally shape inside a @contextmanager mirrors what a class-based __exit__ lets you express, but reads more linearly. The except branch fires only when the body raised, else only when it did not, and finally always. Re-raising with a bare raise is the way to log AND let the caller decide what to do; swallowing here would behave like returning True from __exit__. This is the standard shape for instrumentation: timers, audit logs, span starts and ends.
from contextlib import suppress, ExitStack, contextmanager
# 1) suppress: one-line 'ignore these specific exceptions' block.
with suppress(FileNotFoundError):
open('not-real.txt')
print('survived the missing file')
# 2) ExitStack: dynamically open many context managers in one with-block.
@contextmanager
def tracked(name):
print(f'open {name}')
try:
yield name
finally:
print(f'close {name}')
resources = ['db', 'cache', 'metrics']
with ExitStack() as stack:
handles = [stack.enter_context(tracked(r)) for r in resources]
print('inside, holding:', handles)
# All three close in LIFO order on the way out.contextlib.suppress(*excs) replaces the boilerplate try / except: pass block when you genuinely do not care about a specific failure. ExitStack is the answer to 'I need N context managers but N is computed at runtime': call stack.enter_context(cm) once per resource and the stack guarantees they all close in LIFO order on exit. Together with @contextmanager, these three tools cover almost every 'I want a custom with-block' use case without writing a class at all.