Decorar uma classe inteira

Uma classe também é um objeto

No Python, uma classe também é um objeto, criado quando o interpretador executa o bloco class. Isso significa que você pode passar uma classe para uma função, guardá-la em uma variável e, claro, decorá-la. Um decorador de classe funciona igual ao de função: ele recebe algo (agora a classe), faz o que precisa e devolve algo (a classe, modificada ou não). A sintaxe com arroba é idêntica, só que aplicada em cima de um class em vez de um def.

def com_descricao(cls):
    def descrever(self):
        return f"{cls.__name__} com atributos {list(self.__dict__)}"
    cls.descrever = descrever
    return cls

@com_descricao
class Produto:
    def __init__(self, nome, preco):
        self.nome = nome
        self.preco = preco

p = Produto("Cafe", 20)
print(p.descrever())  # Produto com atributos ['nome', 'preco']

O decorador recebe a classe, acrescenta um método descrever e devolve a classe.

Repare no essencial: com_descricao recebe cls, a classe, define uma função descrever, gruda essa função na classe com cls.descrever = descrever e devolve cls. A partir daí, toda instância de Produto tem o método descrever, mesmo que ele não apareça no corpo da classe. Esse padrão de acrescentar comportamento de fora é útil para aplicar a mesma capacidade a várias classes sem repetir código nem recorrer a herança.

Registro automático de classes

Um uso muito comum de decorador de classe é o registro automático. Imagine um sistema de plugins ou de tipos de relatório: cada nova classe precisa ficar disponível em um catálogo central para ser encontrada pelo nome. Em vez de manter uma lista à mão, que sempre esquecem de atualizar, um decorador cadastra a classe no momento em que ela é definida. O código fica declarativo: basta anotar a classe com @registrar e ela aparece no catálogo.

RELATORIOS = {}

def registrar_relatorio(cls):
    RELATORIOS[cls.tipo] = cls
    return cls

@registrar_relatorio
class RelatorioPDF:
    tipo = "pdf"

@registrar_relatorio
class RelatorioCSV:
    tipo = "csv"

print(RELATORIOS)  # {'pdf': RelatorioPDF, 'csv': RelatorioCSV}
escolhida = RELATORIOS["csv"]
print(escolhida.__name__)  # RelatorioCSV

Cada classe se cadastra sozinha no dicionário RELATORIOS ao ser definida.

O dataclass, exemplo oficial

Você não precisa criar decoradores de classe todo dia, mas encontra um dos mais úteis da linguagem com frequência: o @dataclass. Ele é um decorador de classe da biblioteca padrão que, ao ler os atributos anotados da sua classe, gera automaticamente o __init__, o __repr__ e o __eq__, entre outros. Ou seja, ele modifica a classe acrescentando métodos, exatamente o que você fez à mão nas seções anteriores, só que de forma muito mais completa. Vê-lo como decorador de classe desmistifica a mágica.

from dataclasses import dataclass

@dataclass
class Ponto:
    x: int
    y: int

# O @dataclass gerou __init__, __repr__ e __eq__ automaticamente
p1 = Ponto(1, 2)
p2 = Ponto(1, 2)
print(p1)          # Ponto(x=1, y=2)  <- __repr__ gerado
print(p1 == p2)    # True             <- __eq__ gerado

@dataclass é um decorador de classe que gera métodos a partir dos atributos anotados.

O dataclass tem um módulo inteiro dedicado a ele mais adiante no curso, então aqui basta a ideia: é um decorador de classe. Quando você entende que decorar uma classe é só receber o objeto classe, mexer nele e devolvê-lo, deixa de haver mágica. A mesma técnica que grudou um método descrever no Produto é a que, muito mais elaborada, gera o __init__ do dataclass. Reconhecer o padrão é o que permite ler e escrever esse tipo de código com confiança.