Как ускорить Python с помощью Rust и PyO3

Питон прекрасен своей лаконичностью и богатым выбором сторонних пакетов, но когда дело доходит до тяжелых расчетов, парсинга гигабайтов данных или валидации миллионов объектов, интерпретатор начинает замедляться. Раньше в таких ситуациях приходилось писать расширения C API вручную. Процесс этот трудоемкий: отладка утечек памяти и ручное управление ссылками быстро сводят на нет всю радость от разработки.

Несколько лет назад появился более безопасный путь. Библиотека PyO3 дает возможность писать native-модули на Rust, сохраняя строгие гарантии работы с памятью и понятные абстракции.

Если вы пользовались валидатором Pydantic v2, парсером JSON orjson, фреймворком Polars или токенизатором Tiktoken от OpenAI, вы уже запускали код, созданный с помощью PyO3.

Зачем связывать Rust и Python

Python отлично подходит для бизнес-логики и быстрого прототипирования. Rust выигрывает там, где нужна максимальная скорость, низкое потребление ресурсов и параллелизм.

PyO3 решает две задачи:

Вынос узких мест Python-приложения в отдельные Rust-функции или структуры.
Встраивание интерпретатора Python в Rust-бинарник для исполнения скриптов.

Чаще всего проект используют именно для первого сценария.

Как написать модуль на практике

Для сборки Rust-кода в Python-пакет авторы рекомендуют инструмент maturin. Он отвечает за компиляцию и создание колес wheel.

Создадим простой модуль, который складывает два числа и возвращает результат в виде строки.

Сначала настроим виртуальное окружение:

mkdir string_sum && cd string_sum
python -m venv .env
source .env/bin/activate
pip install maturin
maturin init --bindings pyo3

Логика на стороне Rust в файле src/lib.rs:

use pyo3::prelude::*;

#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

#[pymodule]
mod string_sum {
    use super::*;

    #[pymodule_export]
    fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
        m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
        Ok(())
    }
}

Макросы #[pyfunction] и #[pymodule] автоматически генерируют код для C API Python. Разработчику не нужно вручную преобразовывать типы PyObject в Rust-структуры и обратно.

Запускаем сборку:

maturin develop

Теперь функцию можно вызвать прямо из Python:

import string_sum

result = string_sum.sum_as_string(5, 20)
print(result)  # '25'

Если вы хотите замерить реальную производительность в продакшене, достаточно собрать модуль с флагом оптимизации: maturin develop --release.

Реальные кейсы использования

В Pydantic V2 всю логику проверки типов и валидации переписали на Rust в отдельном пакете pydantic-core. Производительность выросла в разы.

Библиотека Polars вытесняет Pandas при работе с большими наборами данных за счет параллельных вычислений на Rust. Пользователь пишет привычный код на Python, а тяжелая работа уходит в скомпилированный модуль.

В сфере машинного обучения PyO3 применяют Hugging Face в пакете tokenizers и OpenAI в tiktoken. Быстрый разбор текста перед отправкой в нейросеть существенно снижает задержки сервиса.

Вспомогательные инструменты

Вокруг PyO3 сформировался набор полезных утилит:

maturin убирает необходимость писать сложные конфиги для setup.py.
rust-numpy связывает массивы NumPy с Rust без копирования данных в памяти.
pyo3-async-runtimes соединяет asyncio с асинхронным движком Tokio.
pythonize преобразует структуры данных с помощью Serde.

Итоги

Если в проекте есть узкое место по производительности, PyO3 поможет решить проблему без перехода на C или C++. Главная сложность тут упирается в изучение самого Rust. Понимание концепций владения и типов данных потребуется, но рутинную работу с CPython библиотека полностью берет на себя.