Recap do decorador e o functools.wraps

Como o decorador funciona por dentro

No Intermediário você viu que um decorador é uma função que recebe outra função como argumento e devolve uma nova função no lugar dela. A sintaxe com arroba é só um atalho: escrever @registrar em cima de def somar é exatamente o mesmo que escrever somar = registrar(somar) logo depois da definição. Nada de mágico acontece; o nome somar passa a apontar para a função que registrar devolveu. Entender isso com clareza é o que permite avançar para os decoradores mais elaborados deste módulo.

def registrar(funcao):
    def wrapper(*args, **kwargs):
        print(f"chamando {funcao.__name__}")
        return funcao(*args, **kwargs)
    return wrapper

@registrar
def somar(a, b):
    return a + b

# @registrar equivale a: somar = registrar(somar)
print(somar(2, 3))  # chamando somar / 5

O decorador recebe somar, devolve wrapper, e somar passa a ser wrapper.

O wrapper recebe *args e **kwargs para poder envolver qualquer função, com qualquer assinatura. Ele faz o trabalho extra, no caso imprimir uma linha de log, e então chama a função original passando adiante todos os argumentos que recebeu. O return da função original é devolvido sem alteração, para que quem chamou continue recebendo o resultado esperado. Esse esqueleto, recebe função, define wrapper, devolve wrapper, é o mesmo em todo decorador.

O problema silencioso da identidade perdida

Há um efeito colateral que passa despercebido até dar dor de cabeça. Depois de decorar somar, o nome que o Python enxerga para a função não é mais somar, e sim wrapper. Isso acontece porque somar agora aponta para a função interna do decorador, que se chama wrapper. A documentação (a docstring) também some, porque a docstring visível é a do wrapper, que não tem nenhuma. Para código de brincadeira isso não importa, mas ferramentas de depuração, logs, geração de documentação e o próprio pytest usam esses metadados.

def registrar(funcao):
    def wrapper(*args, **kwargs):
        return funcao(*args, **kwargs)
    return wrapper

@registrar
def somar(a, b):
    "Soma dois numeros e devolve o total."
    return a + b

print(somar.__name__)  # wrapper  (deveria ser somar)
print(somar.__doc__)   # None     (a docstring sumiu)

Sem wraps, a função decorada perde o próprio nome e a documentação.

A solução em uma linha: functools.wraps

A biblioteca padrão resolve o problema com functools.wraps. Você o aplica como decorador sobre a função wrapper, passando a função original como argumento. Ele copia os metadados importantes da original para o wrapper: o __name__, a __doc__, o __module__ e outros. O resultado é que a função decorada continua se apresentando com a própria identidade, mesmo estando embrulhada. A partir daqui, todo decorador deste curso usa wraps, e você deveria fazer o mesmo em qualquer código de verdade.

import functools

def registrar(funcao):
    @functools.wraps(funcao)
    def wrapper(*args, **kwargs):
        return funcao(*args, **kwargs)
    return wrapper

@registrar
def somar(a, b):
    "Soma dois numeros e devolve o total."
    return a + b

print(somar.__name__)  # somar
print(somar.__doc__)   # Soma dois numeros e devolve o total.

Uma linha, @functools.wraps(funcao), preserva nome e docstring.