O atalho: @contextlib.contextmanager

Um context manager em uma função

A classe com __enter__ e __exit__ é poderosa, mas exagerada para casos simples. O módulo contextlib traz um atalho elegante: o decorador contextmanager. Você escreve uma função geradora, aquelas com yield, e o decorador a transforma em um context manager. A leitura fica natural: tudo que vem antes do yield é a preparação, equivalente ao __enter__; o valor que você cede com yield é o que vai para o as; e tudo que vem depois do yield é a limpeza, equivalente ao __exit__.

from contextlib import contextmanager

@contextmanager
def conexao(host):
    print(f"conectando a {host}")   # entrada (__enter__)
    conn = {"host": host, "ativa": True}
    try:
        yield conn                   # o valor vai para o as
    finally:
        conn["ativa"] = False
        print(f"desconectando de {host}")  # saida (__exit__)

with conexao("servidor") as c:
    print(f"usando, ativa = {c['ativa']}")
# conectando a servidor
# usando, ativa = True
# desconectando de servidor

Antes do yield é a entrada; depois do yield, dentro do finally, é a limpeza.

Compare com a classe Conexao da aula anterior: o comportamento é idêntico, mas o código cabe em menos linhas e se lê de cima para baixo como uma pequena história. Por isso, na prática do dia a dia, essa é a forma mais comum de criar context managers próprios. A classe fica reservada para quando você precisa guardar bastante estado ou reaproveitar o objeto em contextos diferentes.

Por que a limpeza vai no finally

Há um detalhe que faz toda a diferença: envolver o yield em um try com finally. Quando ocorre uma exceção dentro do bloco with, o Python a repassa para dentro da função geradora, no ponto do yield, como se ela tivesse acontecido ali. Se a limpeza estivesse solta depois do yield, sem o finally, essa exceção pularia a limpeza, exatamente o problema que o with deveria resolver. Colocar a limpeza no finally garante que ela rode aconteça o que acontecer, fiel à promessa do with.

from contextlib import contextmanager

@contextmanager
def recurso():
    print("adquire")
    try:
        yield
    finally:
        print("libera")  # roda mesmo se o bloco lancar excecao

try:
    with recurso():
        print("usando")
        raise RuntimeError("falhou")
except RuntimeError:
    print("erro tratado fora")
# adquire
# usando
# libera        <- a limpeza rodou apesar do erro
# erro tratado fora

O finally garante a limpeza mesmo quando o bloco with lança uma exceção.

Classe ou função: quando usar cada uma

Com as duas formas na mão, surge a pergunta: qual usar? A resposta é sobre estado e reaproveitamento. Se o gerenciador é simples, faz uma preparação e uma limpeza e pronto, a função com yield é mais curta e clara, e deve ser a primeira escolha. Se o gerenciador guarda bastante estado, tem métodos próprios ou precisa ser reutilizado como objeto em vários lugares, a classe com __enter__ e __exit__ organiza melhor. As duas produzem context managers equivalentes; muda a ergonomia.

Uma boa heurística: comece pela função com yield e só migre para a classe se sentir falta de estrutura. O contextlib ainda traz utilitários que combinam com essa abordagem, como o suppress e o ExitStack, que a próxima aula apresenta. O importante é que, tendo entendido o protocolo __enter__ e __exit__ e o atalho com yield, você já consegue criar gerenciadores de contexto para praticamente qualquer recurso com começo e fim.