Escrevendo context managers personalizados em Python

Motivação

Você provavelmente já conhece a instrução with, uma forma concisa de lidar com recursos em Python. Mas já se perguntou como isso funciona por baixo dos panos? O poder da palavra-chave with vem dos context managers.

Context managers são um padrão fundamental em Python que trazem uma abordagem estruturada ao gerenciamento de recursos. Eles garantem que os recursos sejam adquiridos, usados corretamente e, por fim, liberados ou limpos — mesmo quando acontecem erros ou exceções.

Isso é especialmente importante ao lidar com recursos como arquivos, conexões de rede ou acessos a banco de dados. Usando context managers, escrevemos códigos mais limpos e confiáveis, sem a preocupação de esquecer de fechar um arquivo ou liberar um lock.

Neste tutorial, vamos além dos context managers padrão do Python e aprender a escrever os nossos próprios.

Entendendo os context managers em Python

Nos bastidores, context managers são objetos que definem dois métodos especiais: __enter__ e __exit__. O método __enter__ é chamado quando você entra no bloco with, e seu valor de retorno é atribuído a uma variável dentro desse bloco. Já o método __exit__ é chamado quando o bloco with termina, independentemente de concluir normalmente ou lançar uma exceção.

Essa estrutura garante o manuseio adequado de recursos. Vamos ver alguns context managers embutidos do Python para entender melhor:

1. Gerenciamento de arquivos

Considere o exemplo clássico de abrir um arquivo:

with open('file.txt', 'r') as file:
   data = file.read()

Aqui, open('file.txt', 'r') atua como o context manager. Ao entrar no bloco with, o método enter do objeto arquivo é chamado, abrindo o arquivo e atribuindo-o à variável file. Em seguida, você pode usar file.read() para acessar o conteúdo.

Depois, o método __exit__ do objeto arquivo é garantidamente chamado quando o bloco with termina, mesmo se ocorrer uma exceção. Esse método se encarrega de fechar o arquivo, evitando que você deixe descritores de arquivo abertos por aí.

2. Travamento de threads

Indo além de arquivos, context managers também podem ser usados para sincronização de threads com threading.Lock():

import threading

lock = threading.Lock()
with lock:
   # Seção crítica
   print("This code is executed under lock protection.")

Aqui, lock é um objeto threading.Lock, outro context manager. Ao entrar no bloco with, o método __enter__ adquire o lock, garantindo que apenas uma thread execute a seção crítica por vez. Por fim, o método __exit__ libera o lock ao sair do bloco with, permitindo que outras threads prossigam.

3. Conexões com banco de dados

Da mesma forma, context managers podem gerenciar conexões com bancos de dados:

import sqlite3

with sqlite3.connect('database.db') as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM table")
    rows = cursor.fetchall()

A chamada sqlite3.connect('database.db') é um context manager. Entrar no bloco with estabelece a conexão com o banco, atribuindo-a a connection. Você então usa um cursor para interagir com o banco. O método __exit__ garante que a conexão seja fechada ao sair do bloco with, evitando vazamento de recursos.

4. Sockets de rede

Context managers também podem lidar com comunicação de rede:

import socket

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(('localhost', 8080))
    s.sendall(b'Hello, world')
    data = s.recv(1024)

Aqui, socket.socket(socket.AF_INET, socket.SOCK_STREAM) cria um objeto socket que atua como context manager. O __enter__ cria o socket e, dentro do bloco with, você pode conectar, enviar e receber dados. O __exit__ garante que o socket seja fechado corretamente ao final.

5. Varredura de diretórios

Desta vez, os.scandir('.') oferece uma forma de iterar pelas entradas de um diretório:

with os.scandir('.') as entries:
   for entry in entries:
       print(entry.name)

os.scandir('.') atua como context manager aqui. O __enter__ inicia a varredura do diretório e você pode iterar pelas entradas dentro do bloco with. O __exit__ limpa os recursos da varredura ao sair.

Como você viu, a instrução with, impulsionada pelos context managers, simplifica o gerenciamento de recursos ao cuidar automaticamente da alocação e liberação.

Escrevendo context managers personalizados em Python

Agora, vamos ver como escrever seus próprios managers para ter um controle mais fino sobre o gerenciamento de recursos.

Existem duas abordagens principais para criar context managers personalizados: baseadas em classe e baseadas em função.

Abordagem baseada em classe

A abordagem baseada em classe é a mais estruturada e flexível. Aqui, você define uma classe que implementa os métodos especiais __enter__ e __exit__.

Veja um exemplo de uma classe Timer que mede o tempo de execução:

import time


class Timer:
   def __enter__(self):
       self.start_time = time.time()
       return self

   def __exit__(self, exc_type, exc_value, traceback):
       self.end_time = time.time()
       elapsed_time = self.end_time - self.start_time
       print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with Timer() as timer:
       # Code block to measure the execution time
       time.sleep(2)  # Simulate some time-consuming operation
Elapsed time: 2.002082347869873 seconds

Aqui, a classe Timer define o método __enter__ para capturar o horário de início ao entrar no bloco with. Ela retorna self para permitir acesso ao objeto dentro do bloco. O método __exit__ calcula o tempo decorrido ao sair do bloco with e o imprime.

Abordagem baseada em função

Se você prefere algo mais conciso — ainda que menos flexível — a abordagem baseada em função pode ser uma boa opção. Aqui, você usa o decorator contextmanager do módulo contextlib para converter uma função em context manager.

Veja como funciona:

import time
from contextlib import contextmanager


@contextmanager
def timer():
   start_time = time.time()
   yield
   end_time = time.time()
   elapsed_time = end_time - start_time
   print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with timer():
       time.sleep(2)
Elapsed time: 2.0020740032196045 seconds

O decorator @contextmanager transforma a função timer em um context manager. Dentro da função, start_time é capturado e a instrução yield pausa a execução, permitindo que o código dentro do bloco with rode.

Por fim, a funcionalidade do __exit__ é obtida ao capturar o horário de término e imprimir o tempo decorrido.

Na prática, você escreve a lógica do __enter__ antes do yield e a lógica do __exit__ depois dele. As duas abordagens chegam ao mesmo resultado; a escolha depende da sua preferência por estrutura e legibilidade.

Exemplos práticos de context managers

Agora que você sabe escrever context managers personalizados em Python, vamos ver alguns exemplos práticos onde eles são bem úteis.

Gerenciando operações com arquivos

Operações com arquivos são comuns em muitas aplicações. Ler ou escrever da forma correta garante integridade dos dados e eficiência no uso de recursos. Vamos considerar um cenário em que queremos escrever texto em um arquivo usando um context manager personalizado.

class FileManager:
   def __init__(self, filename, mode):
       self.filename = filename
       self.mode = mode
       self.file = None

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()


# Example usage
if __name__ == "__main__":
   with FileManager("example.txt", "w") as file:
       file.write("Hello, world!\n")

Neste exemplo, definimos a classe FileManager para imitar a função open do Python, cuidando das operações de arquivo. Veja o que cada método faz:

• __init__(self, filename, mode): é o construtor da classe. Inicializa o objeto FileManager com o nome do arquivo e o modo em que ele deve ser aberto (por exemplo, leitura, escrita, append).
• __enter__(self): é chamado ao entrar no bloco with. Abre o arquivo indicado por filename no modo indicado por mode. O objeto de arquivo aberto é atribuído a self.file e retornado para uso dentro do bloco with.
• __exit__(self, exc_type, exc_value, traceback): é chamado ao sair do bloco with, independentemente de ter ocorrido uma exceção. Garante que o arquivo seja fechado chamando self.file.close() se self.file não for None. Isso evita vazamento de recursos e assegura a limpeza adequada.

Se você inspecionar o conteúdo de example.txt após executar o código acima, verá o texto escrito:

$ cat example.txt
Hello, world!

A vantagem da classe FileManager é poder implementar qualquer lógica customizada para trabalhar com arquivos. Por exemplo, você pode modificá-la para manter um log dos arquivos abertos e fechados em um arquivo de texto.

import datetime


class FileManager:
   def __init__(self, filename, mode, log_filename="file_log.txt"):
       self.filename = filename
       self.mode = mode
       self.log_filename = log_filename
       self.file = None

   def log_action(self, action):
       timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
       with open(self.log_filename, "a") as log_file:
           log_file.write(f"{timestamp} - {action}: {self.filename}\n")

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       self.log_action("Opened")
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()
           self.log_action("Closed")


# Example Usage
with FileManager("example.txt", "r") as file:
   content = file.read()
   print(content)

Hello, world!

Nesta versão modificada:

• O módulo datetime é importado para trabalhar com carimbos de data e hora.
• O método log_action foi adicionado para registrar ações em arquivos.
• O parâmetro log_filename foi adicionado para definir o nome do arquivo de log. O padrão é 'file_log.txt'.
• O nome do arquivo e a ação (abriu ou fechou) são acrescentados ao log usando log_action dentro dos métodos __enter__ e __exit__.
• Dentro de log_action, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") obtém o timestamp no formato "YYYY-MM-DD HH:MM:SS".
• O timestamp é incluído no log junto com o nome do arquivo e a ação.

Agora, se imprimirmos o conteúdo de file_log.txt, veremos os registros:

$ cat file_log.txt
2024-04-03 12:37:02 - Opened: example.txt
2024-04-03 12:37:02 - Closed: example.txt

Outro caso de uso prático que podemos implementar como context manager personalizado é um gerenciador de banco de dados SQLite:

import sqlite3


class DatabaseConnection:
   def __init__(self, database_name):
       self.database_name = database_name
       self.connection = None

   def __enter__(self):
       self.connection = sqlite3.connect(self.database_name)
       return self.connection

   def __exit__(self, exc_type, exc_value, traceback):
       if self.connection:
           self.connection.close()

Mais uma vez, isso imita a função sqlite3.connect, mas pode ser aprimorado com lógica adicional de gerenciamento do banco. Vamos verificar se o manager está funcionando com algumas funções de exemplo:

# Example usage
def create_table():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           """CREATE TABLE IF NOT EXISTS users (
                           id INTEGER PRIMARY KEY,
                           username TEXT,
                           email TEXT)"""
       )


def insert_data(username, email):
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           "INSERT INTO users (username, email) VALUES (?, ?)", (username, email)
       )
       connection.commit()


def fetch_data():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute("SELECT * FROM users")
       return cursor.fetchall()


if __name__ == "__main__":
   create_table()
   insert_data("john_doe", "john@example.com")
   insert_data("jane_doe", "jane@example.com")
   users = fetch_data()
   print("Users in the database:")
   for user in users:
       print(user)

Users in the database:
(1, 'john_doe', 'john@example.com')
(2, 'jane_doe', 'jane@example.com')
(3, 'john_doe', 'john@example.com')
(4, 'jane_doe', 'jane@example.com')

Ao dividir as operações do banco em funções e usar o manager personalizado que criamos, o código fica muito mais fácil de manter.

Tópicos avançados relacionados a context managers em Python

Nesta seção, vamos tratar de tópicos avançados sobre context managers em Python, incluindo tratamento de erros e aninhamento de context managers.

Tratamento de erros

Tratar erros é essencial em programação, e context managers em Python oferecem um jeito conveniente de lidar com erros durante operações de gerenciamento de recursos.

No exemplo da classe FileManager, o tratamento de erros pode ser implementado dentro do método __exit__ para garantir a limpeza adequada mesmo se ocorrer uma exceção. Por exemplo, se um erro acontecer ao escrever no arquivo (dentro do bloco with), ainda queremos garantir que o arquivo seja fechado para evitar vazamento de recursos. Podemos modificar o método __exit__ assim:

def __exit__(self, exc_type, exc_value, traceback):
   if self.file:
       self.file.close()
       if exc_type is not None:
           print(f"An error occurred: {exc_value}")

Aqui, primeiro fechamos o arquivo com self.file.close() para garantir a limpeza. Em seguida, verificamos se houve exceção (exc_type is not None). Se sim, imprimimos uma mensagem de erro para indicar o que deu errado.

Aqui:

• exc_type é a classe da exceção, como AttributeError e IndexError.
• exc_value é a instância da exceção com a mensagem.
• traceback é o traceback completo do erro.

Vamos executar a classe FileManager com um erro proposital:

with FileManager("example.txt", "r") as file:
   # Simule um erro de exemplo
   content = file.read()
   print(content[150])

An error occurred: string index out of range

Então, para implementar lógica personalizada ao tratar exceções nos seus managers, aqui vai um fluxo de exemplo que você pode seguir:

def __exit__(self, exc_type, exc_value, traceback):
   # Realize a limpeza
   ...

   # Trate a exceção
   if isinstance(exc_value, AttributeError):
       # Trate AttributeError aqui
       print(f"Exception message: {exc_value}")
   # Trate outra exceção
   elif isinstance(exc_value, AnotherError):
       ...

Aninhando context managers

Aninhar context managers permite gerenciar vários recursos de forma simultânea, cada um com seu próprio context manager, dentro de uma única instrução with. Isso é especialmente útil em operações complexas que envolvem múltiplos recursos. Vamos ver como aninhar context managers nos nossos exemplos.

Considere um cenário em que queremos abrir um arquivo e realizar operações em um banco no mesmo contexto. Talvez queiramos inserir o conteúdo de um arquivo CSV como linhas no banco. Podemos fazer isso aninhando os context managers FileManager e DatabaseConnection em uma única instrução with:

if __name__ == "__main__":
   with FileManager("example.txt", "w") as file, DatabaseConnection("example.db") as connection:
       ...

Neste exemplo, os context managers FileManager e DatabaseConnection estão aninhados na mesma instrução with. Isso nos permite abrir o arquivo e estabelecer a conexão com o banco simultaneamente. As operações subsequentes de escrita no arquivo e no banco ocorrem dentro desse contexto combinado.

Um ponto de atenção é a legibilidade do código. Eu recomendo não exceder dois contextos simultâneos.

Conclusão

Context managers em Python oferecem uma abordagem estruturada para gerenciamento de recursos dentro de um with, garantindo alocação e liberação adequadas.

Ao longo deste tutorial, vimos os fundamentos de context managers, passamos por exemplos práticos e discutimos tópicos avançados como tratamento de erros e aninhamento. Ao dominar a escrita de context managers personalizados, você passa a escrever códigos mais limpos e confiáveis, melhora o tratamento de erros e gerencia recursos com eficiência.

O tema de escrever context managers personalizados se encaixa na categoria de "Python avançado". Aqui vão alguns recursos relacionados para você levar suas habilidades ainda mais longe:

• Curso Writing functions in Python
• Trilha de aprendizado Becoming a Python Programmer