Last Updated: September 09, 2019
· maebert

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():

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

with log_runtime('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):

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

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

    def __call__(self, 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"{}: {}".format(self.label, time.time() - self.start))

and use it like this:

def foo():

and like this:

with log_runtime(label="bar"):

2 Responses
Add your response

Thanks a lot!

over 1 year ago ·

This is a great idea, and you can even use the contextlib implementation in 3.2+ (or contextlib2 for Python 2.7):

over 1 year ago ·