Integrating Python with Rust for High-Performance Extensions

[ryanmaynard]

Integrating Python with Rust for High-Performance Extensions

Python is accessible and useful for many things, but sometimes you need more speed. That’s where Rust comes in. By integrating Rust with Python, you can create high-performance extensions that give your Python code a serious speed boost. Let’s explore how to do this using PyO3 and Rust’s Foreign Function Interface (FFI).

Why Rust?

Rust is a systems programming language that offers memory safety and performance. It’s designed to prevent the kinds of bugs that often plague C and C++ code, making it a solid choice for writing safe, fast extensions.

Setting Up Your Environment

First, you need Rust and Cargo installed. If you don’t have them, install Rust from rustup.rs.

Next, you’ll need to create a new Rust project.

cargo new --lib my_rust_extension
cd my_rust_extension

Add `pyo3` to your `Cargo.toml` file.

[dependencies]
pyo3 = { version = "0.15", features = ["extension-module"] }

Writing Your Rust Code

Create a simple function in Rust that you want to call from Python. Edit the `src/lib.rs` file:

use pyo3::prelude::*;

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

#[pymodule]
fn my_rust_extension(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

This code defines a function `sum_as_string` that takes two integers, adds them, and returns the result as a string. The `#[pymodule]` attribute creates a Python module.

Building and Installing the Extension

Next, build the extension. Create a `setup.py` file in your project root:

from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    name="my_rust_extension",
    version="0.1.0",
    rust_extensions=[RustExtension("my_rust_extension", binding=Binding.PyO3)],
    packages=["my_rust_extension"],
    zip_safe=False,
)

Install the `setuptools-rust` package:

pip install setuptools-rust

Now, build and install your extension:

python setup.py install

Using Your Rust Extension in Python

After installation, you can use the Rust extension in Python like any other module.

import my_rust_extension

result = my_rust_extension.sum_as_string(5, 10)
print(result)  # Output: "15"

Performance Comparison

To see the performance improvements, let’s compare this Rust function to a pure Python equivalent.

Python Version:

def sum_as_string(a, b):
    return str(a + b)

import time

start = time.time()
for _ in range(1000000):
    sum_as_string(5, 10)
print("Python:", time.time() - start)

Rust Version:

import my_rust_extension
import time

start = time.time()
for _ in range(1000000):
    my_rust_extension.sum_as_string(5, 10)
print("Rust:", time.time() - start)

In most cases, the Rust version will be significantly faster due to Rust’s performance characteristics.

Conclusion

Integrating Rust with Python using PyO3 allows you to create high-performance extensions that can drastically speed up your Python code. By combining the simplicity and ease of Python with the speed and safety of Rust, you can tackle performance-critical tasks more effectively. Try integrating Rust into your next project and see the difference it makes. As a closing note, while I have plenty of Python experience, I'm fairly new to Rust, so if readers have any suggestions, corrections, or additional helpful information - please do share it.

 

[brb3]

This is a great way to creep Rust into a code base where it makes sense to get the performance benefits.
Additionally, I know you can do the same with Rust + Ruby, and probably others as well. (I'll let a Rubyist around here expand on that).

I'm a strong proponent of "right tool for the job" so a case like this where you can move rapidly with Python and then reduce bottlenecks later with Rust is just :chef-kiss: