Context Manager + Decorator

(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 3 it's trivial, you just make a normal context manager and inherit from contextlib.ContextDecorator:

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!
>>>
>>> @wrapper
... def foo(text):
...     print(text)
...
...
>>> foo("I got decorated")
Entering!
I got decorated
Exiting!
>>>

Old Python 2 version

In Python 2, it was a bit more work - I originally wrote this blog post because it was hard for me to remember. You just do the decoration in __call__:

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

Tags