paint-brush
Calling Rust from Python with pyo3by@nfrankel
1,614 reads
1,614 reads

Calling Rust from Python with pyo3

by Nicolas FränkelNovember 2nd, 2023
Read on Terminal Reader
Read this story w/o Javascript
tldt arrow

Too Long; Didn't Read

This post discusses the transition from a low-level integration approach to utilizing the pyo3 library for Rust and Python integration. pyo3 simplifies the process, handling Python types and providing tools like maturin. The tutorial outlines the migration steps, including code adjustments and building Python packages for seamless integration. The enhanced integration enables you to call Rust functions from Python with ease, improving development efficiency.
featured image - Calling Rust from Python with pyo3
Nicolas Fränkel HackerNoon profile picture


I got plenty of feedback on my post about Calling Rust from Python:



Many comments mentioned pyo3, and that I should use it instead of cooking my own. Thanks to the authors, I checked: in this post, I explain what it is and how I migrated my code.


What is pyo3?

Rust bindings for Python, including tools for creating native Python extension modules. Running and interacting with Python code from a Rust binary is also supported.


— PyO3 user guide [1]


Indeed, pyo3 fits my use case, calling Rust from Python. Even better, it handles converting Python types to Rust types and back again. Finally, it offers the maturin utility to make the interaction between the Python project and the Rust project seamless.


Maturin

Build and publish crates with pyo3, rust-cpython, cffi and uniffi bindings as well as rust binaries as python packages.


— Maturin on GitHub [2]


maturin is available via pip install. It offers several commands:


  • new: create a new Cargo project with maturin configured
  • build: build the wheels and store them locally
  • publish: build the crate into a Python package and publish it to pypi
  • develop: build the crate as a Python module directly into the current virtual environment, making it available to Python


Note that Maturin started as a companion project to pyo3 but now offers rust-cpython, cffi and uniffi bindings.


Migrating the project

The term migrating is a bit misleading here since we will start from scratch to fit Maturin’s usage. However, we will achieve the same end state. I won’t paraphrase the tutorial since it works seamlessly. Ultimately, we have a fully functional Rust project with a single sum_as_string() function, which we can call in a Python shell. Note the dependency to pyo3:


pyo3 = "0.20.0"


The second step is to re-use the material from the previous project. First, we add our compute() function at the end of the lib.rs file:


#[pyfunction]                                                                            //1
fn compute(command: &str, a: Complex<f64>, b: Complex<f64>) -> PyResult<Complex<f64>> {  //2-3
    match command {
        "add" => Ok(a + b),
        "sub" => Ok(a - b),
        "mul" => Ok(a * b),
        _ => Err(PyValueError::new_err("Unknown command")),                              //4
    }
}


  1. The pyfunction macro allows the use of the function in Python
  2. Use regular Rust types for parameters; pyo3 can convert them
  3. We need to return a PyResult type, which is an alias over Result<T, PyErr>
  4. Return a specific Python error if the command doesn’t match


pyo3 automatically handles conversion for most types. However, complex numbers require an additional feature. We also need to migrate from the num crate to the num-complex:


pyo3 = { version = "0.20.0" , features = ["num-complex"]}
num-complex = "0.4.4"


To convert custom types, you must implement traits FromPyObject for parameters and ToPyObject for return values.


Finally, we only need to add the function to the module:


#[pymodule]
fn rust_over_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    m.add_function(wrap_pyfunction!(compute, m)?)?;              //1
    Ok(())
}
  1. Add the function to the module


At this point, we can use Maturin to test the project:


maturin develop


After the compilation finishes, we can start a Python shell in the virtual environment:


python

>>> from rust_over_pyo3 import compute
>>> compute('add',1+3j,-5j)
(1-2j)
>>> compute('sub',1+3j,-5j)
(1+8j)


Finishing touch

The above setup allows us to use Rust from a Python shell but not in a Python file. To leverage the default, we must create a Python project inside the Rust project, whose name matches the Rust module name. Since I named my lib rust_over_pyo3, here’s the overall structure:


my-project
├── Cargo.toml
├── rust_over_pyo3
│   └── main.py
├── pyproject.toml
└── src
    └── lib.rs


To use the Rust library in Python, we need first to build the library.


maturin build --release


We manually move the artifact from /target/release/maturin/librust_over_pyo3.dylib to rust_over_pyo3.so under the Python package. We can also run cargo build --release instead; in this case, the source file is directly under /target/release.


At this point, we can use the library as any other Python module:


from typing import Optional
from click import command, option

from rust_over_pyo3 import compute                                                (1)

@command()
@option('--add', 'command', flag_value='add')
@option('--sub', 'command', flag_value='sub')
@option('--mul', 'command', flag_value='mul')
@option('--arg1', help='First complex number in the form x+yj')
@option('--arg2', help='Second complex number in the form x\'+y\'j')
def cli(command: Optional[str], arg1: Optional[str], arg2: Optional[str]) -> None:
    n1: complex = complex(arg1)
    n2: complex = complex(arg2)
    result: complex = compute(command, n1, n2)                                    (2)
    print(result)


if __name__ == '__main__':
    cli()


  1. Regular Python import
  2. Look, ma, it works!


Conclusion

In this post, I improved the low-level integration with ctypes to the generic ready-to-use pyo3 library. I barely scratched the surface, though; pyo3 is a powerful, well-maintained library with plenty of features.


I want to thank everyone who pointed me in this direction.


The complete source code for this post can be found on GitHub.


To go further:

[1] PyO3 user guide

[2] maturin


Originally published at A Java Geek on October 29th, 2023