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"

Comments

comments powered by Disqus