Como Acelerar Python com Rust e PyO3

Python é lindo em sua concisão e vasta seleção de pacotes de terceiros, mas quando se trata de cálculos pesados, analisar gigabytes de dados ou validar milhões de objetos, o interpretador começa a ficar lento. Anteriormente, em tais situações, os desenvolvedores precisavam escrever extensões da API C manualmente. Esse processo é trabalhoso: depurar vazamentos de memória e gerenciamento manual de referências rapidamente eliminam toda a alegria do desenvolvimento.

Um caminho mais seguro surgiu há alguns anos. A biblioteca PyO3 permite escrever módulos nativos em Rust mantendo garantias rigorosas de segurança de memória e abstrações claras.

Se você já usou o validador do Pydantic v2, o parser JSON orjson, o framework Polars ou o tokenizador Tiktoken da OpenAI, você já executou código criado com PyO3.

Por que Vincular Rust e Python

Python é excelente em lógica de negócios e prototipagem rápida. Rust vence onde máxima velocidade, baixo consumo de recursos e concorrência são necessários.

PyO3 resolve dois problemas:

1. Transferir gargalos de aplicações Python para funções ou estruturas Rust separadas.
2. Incorporar o interpretador Python em um binário Rust para execução de scripts.

Na maioria das vezes, o projeto é usado para o primeiro cenário.

Como Escrever um Módulo na Prática

Para compilar código Rust em um pacote Python, os autores recomendam a ferramenta 4. Ela lida com a compilação e cria arquivos wheel.

Vamos criar um módulo simples que soma dois números e retorna o resultado como uma string.

Primeiro, vamos configurar o ambiente virtual:

O código Rust no arquivo 5:

As macros 6 e 7 geram automaticamente código para a API C do Python. Os desenvolvedores não precisam converter manualmente tipos 8 para estruturas Rust e vice-versa.

Vamos iniciar a compilação:

Agora a função pode ser chamada diretamente do Python:

Se você quiser medir o desempenho real em produção, basta compilar o módulo com a flag de otimização: 9.

Casos de Uso Reais

No Pydantic V2, toda a lógica de verificação de tipos e validação foi reescrita em Rust em um pacote separado 10. O desempenho aumentou drasticamente.

A biblioteca Polars está substituindo o Pandas ao trabalhar com grandes conjuntos de dados graças a cálculos paralelos em Rust. O usuário escreve código Python familiar, enquanto o trabalho pesado fica no módulo compilado.

Em aprendizado de máquina, PyO3 é usado pela Hugging Face no pacote 11 e pela OpenAI no 12. A análise rápida de texto antes do envio para a rede neural reduz significativamente a latência do serviço.

Ferramentas de Suporte

Um conjunto de utilitários úteis se formou em torno do PyO3:

• 13 elimina a necessidade de escrever configurações complexas para setup.py.
• 14 vincula arrays NumPy com Rust sem copiar dados na memória.
• 15 conecta asyncio com o runtime assíncrono Tokio.
• 16 transforma estruturas de dados usando Serde.

Conclusão

Se o seu projeto tem um gargalo de desempenho, PyO3 ajudará a resolver o problema sem trocar para C ou C++. A principal dificuldade aqui está em aprender Rust. Será necessário entender conceitos de ownership e tipos de dados, mas a biblioteca cuida totalmente do trabalho rotineiro com o CPython.