Python Decorators vs. Context Managers: Have your cake and eat it!

Recently I wrote a decorator that does this:

@log_runtime('calling foo')
def foo():
    do_stuff()

Then, deep in some other function, I wanted to do this:

with log_runtime('do other stuff'):
    do_other_stuff()

That lead to an important realisations: Conceptually, most real decorators are context managers.

I love context managers, so I started writing all decorators as hypercharged hybrid decorator-context-managers:

from functools import wraps

class ContextDecorator(object):
    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

    def __enter__(self):
        # Note: Returning self means that in "with ... as x", x will be self
        return self

    def __exit__(self, typ, val, traceback):
        pass

    def __call__(self, f):
        @wraps(f)
        def wrapper(*args, **kw):
            with self:
                return f(*args, **kw)
        return wrapper

Now, I can write log_runtime like this:

class log_runtime(ContextDecorator):
    def __enter__(self):
        self.start_time = time.time()
        return self
    def __exit__(self, typ, val, traceback):
        # Note: typ, val and traceback will only be not None
        # If an exception occured
        logger.info("{}: {}".format(self.label, time.time() - self.start))

and use it like this:

@log_runtime(label="foo")
def foo():
    do_stuff()

and like this:

with log_runtime(label="bar"):
    do_bar_stuff()