Libraries Run Rust Inside Python (With PyO3)

Pydantic v2, a popular data validation library for Python, relies on a Rust extension called PyO3. This post demonstrates how to create a JSON parser in Rust and expose it to Python using PyO3. The process involves several key steps.

To get Rust code into Python, you use two Rust macros: #[pyfunction] and #[pymodule]. These act like Python decorators, allowing Python to call the Rust functions and handling type conversions and reference counting at the boundary. Maturin, a tool, compiles the Rust crate into a shared library (.so, .dylib, .dll) and integrates it into your Python environment.

The parser returns a JSON tree, which is implemented as a Rust enum. This enum can hold one of several shapes, each carrying data, making it a perfect fit for representing a JSON tree. The tree itself resides entirely in Rust, with Python never directly interacting with it. The PyO3 layer serves as a thin adapter between Rust and Python.

The conversion process from Rust to Python is not free. It involves creating Python objects for nodes in the tree, which can be more work than the parsing itself, especially for large documents. The .into_pyobject function in PyO3 walks the entire JsonValue tree, rebuilding it as native Python objects. This includes creating a dict for objects, a list for arrays, and a float or string for leaf nodes.

When parsing fails, it raises a typed Rust error, which Python interprets as an exception. An implementation of the From trait allows this conversion, turning a parsing error into a Python ValueError carrying the offset where parsing broke. The file-reading path also benefits from this conversion, as std::io::Error already converts to the appropriate Python exception, like FileNotFoundError.

The conversion cost is most significant when the Rust function returns a large structure. To optimize this, you can preallocate a PyDict at the margins, but the bigger win is to avoid materializing the entire tree if the caller won't touch all of it. Instead, return a lazy, Rust-backed view and build Python objects on demand.

When using PyO3, it's crucial to profile the boundary between Rust and Python, not just the algorithm itself. Getting Rust to run fast is the easy part, but the real challenge lies in optimizing the conversion process from Rust values to Python objects.