(Updated 2018 with a more modern Python 3 tip)
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.
In Python 2, you just do the decoration in
from functools import wraps class DecoratorOrContextManager(object): def __enter__(self): print("Entering!") return self def __exit__(self, typ, value, tb): print("Exiting!") return False 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
In Python 3 it's even easier, you just make a normal context manager and inherit
from contextlib import ContextDecorator class DecoratorOrContextManager(ContextDecorator): def __enter__(self): print("Entering!") return self def __exit__(self, typ, value, tb): print("Exiting!") return False
Now you can use it either way, like so:
>>> wrapper = DecoratorOrContextManager() >>> with wrapper: ... print("my context is managed") ... Entering! my context is managed Exiting! >>> >>> @thingy ... def foo(): ... print("I got decorated") ... Entering! I got decorated Exiting! >>>