Python Decorators vs. Context Managers: Pegue o seu bolo e coma-o!

Recentemente, escrevi um decorador que faz isso:

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

Então, profundamente em alguma outra função, eu queria fazer isso:

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

Isso levou a uma importante realização: conceitualmente, a maioria dos decoradores reais são gerenciadores de contexto .

Eu amo gerentes de contexto, então comecei a escrever todos os decoradores como gerentes de contexto-decorador híbridos hipercarregados:

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

Agora, posso escrever log_runtimeassim:

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

e use-o assim:

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

e assim:

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