O protocolo __enter__ e __exit__

Os dois métodos que definem um context manager

O with não é mágico: ele apenas chama dois métodos especiais do objeto que você passa. Ao entrar no bloco, o Python chama o __enter__ do objeto; o valor que esse método devolve é atribuído à variável depois do as. Ao sair do bloco, por qualquer motivo, o Python chama o __exit__, onde vai a limpeza. Qualquer objeto que tenha esses dois métodos pode ser usado com with. Criar um context manager é, no fundo, escrever essas duas funções.

class Conexao:
    def __init__(self, host):
        self.host = host

    def __enter__(self):
        print(f"conectando a {self.host}")
        self.ativa = True
        return self          # vai para a variavel do as

    def __exit__(self, exc_tipo, exc_valor, traceback):
        print(f"desconectando de {self.host}")
        self.ativa = False
        return False         # nao suprime excecoes

with Conexao("servidor") as conn:
    print(f"usando conexao, ativa = {conn.ativa}")
# conectando a servidor
# usando conexao, ativa = True
# desconectando de servidor

Uma classe com __enter__ e __exit__ vira um context manager utilizável com with.

Acompanhe a ordem. O with Conexao("servidor") as conn cria o objeto e chama seu __enter__, que imprime a conexão e devolve self; por isso conn passa a ser a própria conexão. O bloco roda. Na saída, o Python chama __exit__, que imprime a desconexão. Repare que __enter__ costuma devolver self quando o objeto útil é ele mesmo, mas poderia devolver qualquer coisa, um arquivo, um cursor, um valor. É esse retorno que aparece depois do as.

Os parâmetros do __exit__

O __enter__ é simples; o __exit__ tem um detalhe importante. Ele recebe três parâmetros que descrevem a exceção que porventura ocorreu dentro do bloco: o tipo da exceção, o valor (a instância do erro) e o traceback. Se o bloco terminou sem erro, os três chegam como None. Esse desenho permite que o __exit__ saiba se a saída foi tranquila ou por falha, e reaja de formas diferentes, por exemplo desfazendo uma transação em caso de erro e confirmando em caso de sucesso.

class Transacao:
    def __enter__(self):
        print("iniciando transacao")
        return self

    def __exit__(self, exc_tipo, exc_valor, traceback):
        if exc_tipo is None:
            print("commit: tudo certo")
        else:
            print(f"rollback: erro {exc_tipo.__name__}")
        return False

with Transacao():
    print("gravando dados")
    raise ValueError("dado invalido")
# iniciando transacao
# gravando dados
# rollback: erro ValueError
# ... e a excecao continua subindo

O __exit__ inspeciona exc_tipo para decidir entre confirmar e desfazer a operação.

Montando o seu, do começo ao fim

Junte as peças e você tem uma receita: no __init__ guarde a configuração; no __enter__ adquira o recurso e devolva o que o usuário vai manipular; no __exit__ libere o recurso e decida o que fazer com uma eventual exceção. Essa estrutura funciona para arquivos, conexões, bloqueios, medições de tempo, mudanças temporárias de configuração e muito mais. Sempre que você pensar em algo que tem começo e fim garantidos, um context manager de classe é a ferramenta.

Vale separar bem o __init__ do __enter__. O __init__ apenas prepara os dados; a aquisição de verdade acontece no __enter__, quando o with começa. Essa separação permite, por exemplo, criar o objeto uma vez e usá-lo em vários blocos with. Nas próximas aulas você vai ver que, para casos simples, existe um atalho que evita escrever a classe inteira. Mas entender o protocolo de classe é a base: o atalho é apenas uma forma mais curta de produzir os mesmos __enter__ e __exit__.