Decorador que recebe argumentos

Por que precisa de um nível a mais

Um decorador comum recebe a função e devolve o wrapper: dois níveis. Mas repare na diferença entre @registrar e @repetir(vezes=3). No segundo caso há parênteses e argumentos. Isso significa que @repetir(vezes=3) primeiro executa repetir(vezes=3), e o resultado dessa chamada é que vai decorar a função. Ou seja, repetir(vezes=3) precisa devolver um decorador. É esse nível extra que confunde no começo: a função mais de fora não recebe a função a decorar, ela recebe a configuração.

# Decorador SEM argumentos: dois niveis
@registrar
def somar(a, b):
    return a + b
# equivale a: somar = registrar(somar)

# Decorador COM argumentos: tres niveis
@repetir(vezes=3)
def ola():
    print("ola")
# equivale a: ola = repetir(vezes=3)(ola)
#            \_______/ \__/
#             fabrica    funcao a decorar

Com argumentos, a fábrica é chamada primeiro e devolve o decorador que envolve a função.

Guarde a linha de baixo do exemplo: ola = repetir(vezes=3)(ola). Ela mostra que há duas chamadas em sequência. A primeira, repetir(vezes=3), roda a fábrica e devolve um decorador. A segunda, (ola), aplica esse decorador à função ola. Se você consegue ler essa linha com naturalidade, o resto é só montar as três funções aninhadas na ordem certa.

Montando a fábrica de três níveis

Agora o código completo. São três funções, uma dentro da outra. A mais externa, repetir, recebe a configuração (vezes) e existe só para devolver o decorador. A do meio, decorador, é um decorador normal: recebe a função a envolver. A mais interna, wrapper, é onde o trabalho acontece, e ela enxerga tanto vezes (da fábrica) quanto funcao (do decorador), graças à closure. Leia de fora para dentro e depois de dentro para fora, até a lógica assentar.

import functools

def repetir(vezes):
    def decorador(funcao):
        @functools.wraps(funcao)
        def wrapper(*args, **kwargs):
            resultado = None
            for _ in range(vezes):
                resultado = funcao(*args, **kwargs)
            return resultado
        return wrapper
    return decorador

@repetir(vezes=3)
def cumprimentar(nome):
    print(f"Ola, {nome}!")
    return nome

cumprimentar("Ana")
# Ola, Ana!
# Ola, Ana!
# Ola, Ana!

Três níveis: repetir (configuração), decorador (recebe a função), wrapper (executa).

Detalhes que valem lembrar

Dois pontos merecem atenção. Primeiro, o functools.wraps continua indo no wrapper, o nível mais interno, exatamente como no decorador simples. Segundo, existe uma pegadinha clássica: escrever @repetir sem parênteses quando o decorador espera argumentos. Nesse caso, o Python passa a própria função como se fosse o parâmetro vezes, e tudo quebra de um jeito confuso. Se o decorador foi feito para receber configuração, ele sempre é usado com parênteses, mesmo que você aceite os valores padrão: @repetir().

Vale um último enquadramento. A fábrica não é um truque à parte, é o mesmo decorador de sempre com uma camada de configuração por fora. Se um dia você travar, escreva a linha equivalente sem a arroba, como ola = repetir(vezes=3)(ola), e siga as chamadas uma a uma. O padrão de três níveis aparece o tempo todo em bibliotecas reais, de frameworks web a ferramentas de cache, então reconhecê-lo de imediato economiza muito esforço.