Decorators em Python, explicados

Um decorator é uma função que recebe uma função e devolve um substituto para ela. A linha com @ é um atalho para reatribuir o nome.

Depois que você escreve um na mão, a sintaxe para de parecer mágica.

Escreva um sem o @

O wrapper de dentro aceita qualquer coisa, chama a função original e altera o resultado.

# a decorator is a function that returns a function
def shout(fn):
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs).upper()
    return wrapper

def greet(name):
    return f'hello {name}'

loud = shout(greet)
print(loud('ada'))

Ele imprime:

HELLO ADA

Nada de especial aconteceu ainda. Uma função entrou e uma função diferente saiu.

A linha com @ é exatamente isso

@shout acima de uma definição significa greet2 = shout(greet2).

# the @ syntax is the same thing
@shout
def greet2(name):
    return f'hello {name}'

print(greet2('ada'))

Ele imprime:

HELLO ADA

Mesmo resultado. O decorator roda uma vez, no momento da definição.

O wrapper esconde a função original

O nome e a docstring agora são os do wrapper.

# the wrapper hides the original function
print(greet2.__name__, '|', greet2.__doc__)

Ele imprime:

wrapper | None

Isso quebra a saída do help, os depuradores e qualquer coisa que leia __name__.

functools.wraps resolve

@functools.wraps(fn) copia o nome, a docstring e mais alguns atributos para o wrapper.

# functools.wraps keeps the metadata
import functools

def shout_fixed(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        return fn(*args, **kwargs).upper()
    return wrapper

@shout_fixed
def greet3(name):
    """Say hello."""
    return f'hello {name}'

print(greet3.__name__, '|', greet3.__doc__)
print(greet3('ada'))

Ele imprime:

greet3 | Say hello.
HELLO ADA

Coloque isso em todo decorator que você escrever. Não existe caso em que você prefira os metadados do wrapper.

Um decorator que recebe argumentos

@repeat(3) significa “chame repeat(3) e use o que voltar como decorator”. Isso exige mais um nível de aninhamento.

# a decorator that takes arguments needs one more layer
def repeat(times):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            return [fn(*args, **kwargs) for _ in range(times)]
        return wrapper
    return decorator

@repeat(3)
def roll():
    return 4

print(roll())

Ele imprime:

[4, 4, 4]

Três camadas: a que recebe o argumento, o decorator e o wrapper.

Guardando estado

O wrapper é um closure, então ele consegue guardar estado entre chamadas. Pendurar esse estado no próprio wrapper deixa ele legível de fora.

# state in the closure, such as counting calls
def counted(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        wrapper.calls += 1
        return fn(*args, **kwargs)
    wrapper.calls = 0
    return wrapper

@counted
def work(n):
    return n * 2

work(1); work(2); work(3)
print('called', work.calls, 'times')

Ele imprime:

called 3 times

O que já vem na biblioteca padrão

functools.lru_cache é um decorator que guarda resultados por argumento. Num Fibonacci recursivo ingênuo, ele transforma trabalho exponencial em linear.

# caching is a decorator in the standard library
@functools.lru_cache(maxsize=None)
def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

print(fib(30), fib.cache_info())

Ele imprime:

832040 CacheInfo(hits=28, misses=31, maxsize=None, currsize=31)

As informações do cache mostram quantas chamadas foram respondidas sem executar o corpo da função.

O que lembrar

• @decorator significa name = decorator(name).

• O decorator roda no momento da definição, o wrapper roda no momento da chamada.

• Use sempre functools.wraps, ou você perde o nome e a docstring.

• Um decorator com argumentos precisa de três funções aninhadas.

A maioria dos decorators que você vai escrever é log, medição de tempo, cache, retentativa ou verificação de acesso. Todos têm o mesmo formato do primeiro exemplo desta página.