Integrating Rust into Python via PyO3

Whenever you utilize Pydantic v2 for data validation, you aren't just running Python; you are triggering a Rust extension. The heavy lifting is handled by pydantic-core, which is engineered using PyO3.

This article explores how to build a similar bridge. We will examine a JSON parser written in Rust and exposed to Python, allowing it to be imported as a standard package. Crucially, we will analyze the "final mile"—converting Rust results into Python objects—which can often be more computationally expensive than the actual parsing.

The Pipeline: From Rust to import

Getting Rust code to execute within a Python environment follows a specific workflow:

Implementation Checklist

• Create a standard Rust module.
• Use #[pyfunction] and #[pymodule] macros to handle the "wiring."
• Let Maturin compile the crate into a shared library.
• Install the library into the Python virtual environment.

The macros act similarly to Python decorators; they rewrite the underlying function to manage type conversion and reference counting at the language boundary. You don't need to manually write C-API glue code.

The Rust-Side Logic

Before the data reaches Python, it exists as a Rust value. In a specialized cohort, students build a recursive-descent JSON parser from scratch (without serde). The results are impressive:

• Josh's version: Outperformed CPython's C-based json module on real-world data.
• Jochen's version: Ran up to $3.5\times$ faster than the native Python implementation.

Data Representation

To represent a JSON tree, a Rust enum is used because it can hold multiple shapes:

pub enum JsonValue {
    Null,
    Boolean(bool),
    Number(f64),
    String(String),
    Array(Vec<JsonValue>),
    Object(HashMap<String, JsonValue>),
}

Exposing the Functionality

To make the parse_json function visible to Python, the following signature is used:

#[pyfunction]
fn parse_json<'py>(py: Python<'py>, input: str) -> PyResult<Bound<'py, PyAny>> {
    parse(input)?.into_pyobject(py)
}

Breaking Down the Signature

The actual parsing is done by parse(input), but the magic (and the cost) happens at .into_pyobject(py).

The "Materialization Tax"

The conversion from Rust to Python is not free. The .into_pyobject method must traverse the entire JsonValue tree and instantiate native Python objects.

Warning: For a large document, the process of creating Python objects can dominate the total execution time, potentially costing more than the parsing itself.

This is achieved by implementing the IntoPyObject trait:

impl<'py> IntoPyObject<'py> for JsonValue {
    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
        match self {
            JsonValue::Null => Ok(py.into_any()),
            JsonValue::Object(obj) => {
                let py_dict = PyDict::new(py);
                for (k, v) in obj {
                    // Convert keys and values into Python objects
                    py_dict.set_item(k.into_pyobject(py)?, v.into_pyobject(py)?)?;
                }
                Ok(py_dict.into_any())
            }
            // ... handling for arrays, strings, and booleans
        }
    }
}

If a document contains $100,000$ values, the boundary must create $\approx 100,000$ Python objects. The time complexity of this materialization is $O(n)$, where $n$ is the number of nodes in the tree.

Handling Errors Across the Boundary

Rust's typed errors must be translated into Python exceptions. This is handled via the From trait:

impl From<JsonError> for PyErr {
    fn from(err: JsonError) -> PyErr {
        match err {
            JsonError::UnterminatedString { position } => 
                PyValueError::new_err(format!("Unterminated string at position {position}")),
            // ... other error variants
        }
    }
}

Because of this, a JsonError becomes a ValueError in Python. Similarly, std::io::Error automatically maps to FileNotFoundError, ensuring the Rust layer remains invisible to the end user.

Strategic Takeaways for Porting Code

When moving Python logic to Rust via PyO3, consider the return type:

1. Scalars: If the function returns a single value (e.g., an integer or boolean), port it immediately.
2. Large Structures: If the function returns a massive tree or list, the conversion is your primary bottleneck.

Optimization Tips:

• Preallocating PyDict can provide marginal gains.
• Architectural Shift: Avoid materializing the entire tree. Instead, return a lazy, Rust-backed view and only create Python objects for the specific elements the caller accesses.

Final Rule: Always profile the boundary, not just the algorithm.