I always forget how to write a class that can be used as both a decorator and a context manager. It's a handy trick, but I don't do it often enough to remember, so, here it is:
from functools import wraps class DecoratorOrContextManager(object): def __init__(self): # Add more args & state as needed def __call__(self, f): # For use as a decorator. @wraps(f) def wrapper(*args, **kw): # By using the context manager internally, we ensure that the # cleanup in __exit__ happens even if f() raises an exception. with self: return f(*args, **kw) return wrapper def __enter__(self): print "Setup code goes here" return self def __exit__(self, typ, value, tb): print "Cleanup code goes here"
Now you can use the context manager like so:
with DecoratorOrContextManager() as mgr: print "my context is managed"
Or the decorator like so:
@DecoratorOrContextManager() def foo(): print "I got decorated"
Or slightly nicer, no trailing parens:
def nicer_decorator(f, *decorator_args, **decorator_kwargs): # This is a convenience for slightly nicer decorator syntax. @wraps(f) def wrapper(*args, **kwargs): with DecoratorOrContextManager(*decorator_args, **decorator_kwargs): return f(*args, **kwargs) return wrapper @nicer_decorator def foo2(): print "I got decorated too"