Performance Leap: How to Build High-Performance Python Modules with Rust

Introduction

Python, as a flexible and easy-to-use programming language, is loved by developers. However, when dealing with compute-intensive tasks, Python’s performance bottlenecks can often be frustrating. At this point, the Rust language, with its excellent performance and memory safety features, becomes an ideal choice for extending Python.

With tools like PyO3 and maturin, we can easily build Python modules using Rust, achieving a complementary advantage of both languages: the ease of use of Python and the high performance of Rust. This article will guide you step by step from scratch on how to create a Python module written in Rust.

Why Use Rust to Extend Python?

There are several significant advantages to combining Rust with Python:

Performance Improvement: In numerical computation tasks, Rust can be 10-100 times faster than Python.
Memory Safety: Rust’s compiler can prevent common memory errors at compile time.
Interoperability: Seamless integration with existing Python codebases.
Ecology: Access to Rust’s ever-growing crate ecosystem.

Preparation

Before we begin, you need to install the following tools:

Rust (install via rustup)
Python 3.8+
pip (latest version recommended)
venv (for creating virtual environments)
maturin (Python/Rust build tool)

Install maturin via pip:

pip install maturin

Creating a Project

Create a new Rust library using maturin, which will become your Python module:

maturin new rust_python_module
cd rust_python_module

This will create a project with the following structure:

rust_python_module/
├── Cargo.toml
├── .gitignore
└── src/
    └── lib.rs

Writing Rust Code

Open src/lib.rs and replace its content:

use pyo3::prelude::*;

/// Format the sum of two numbers as a string
#[pyfunction]
fn add(a: usize, b: usize) -> usize {
    // Simply return the sum of two numbers
    a + b
}

/// Python module implemented in Rust
#[pymodule]
fn rust_python_module(m: &mut PyModule) -> PyResult<()> {
    // Add Rust function to Python module
    m.add_function(wrap_pyfunction!(add, m)?)?;
    Ok(())
}

This code defines a simple add function that takes two integers and returns their sum. The #[pyfunction] annotation makes this function callable from Python, while the #[pymodule] defines the Python module.

Building the Python Module

First, create a new virtual environment and activate it:

python3 -m venv venv
source venv/bin/activate  # macOS/Linux
venv\Scripts\activate     # Windows

Install maturin in the virtual environment:

pip install maturin

Then run:

maturin develop

This will compile the Rust code and install it into your currently activated virtual environment.

Testing in Python

Create a Python file and import the module:

import rust_python_module

# Call the function we implemented in Rust
result = rust_python_module.add(2, 3)
print(f"2 + 3 = {result}")  # Output: 2 + 3 = 5

Run the Python script:

python main.py

Advanced Example: Implementing the Fibonacci Sequence

Let’s add a more complex example to showcase Rust’s performance advantages. Modify src/lib.rs:

use pyo3::prelude::*;

/// Format the sum of two numbers as a string
#[pyfunction]
fn add(a: usize, b: usize) -> usize {
    a + b
}

/// Calculate the nth Fibonacci number
#[pyfunction]
fn fibonacci(n: usize) -> usize {
    match n {
        0 => 0,
        1 => 1,
        _ => fibonacci(n - 1) + fibonacci(n - 2),
    }
}

/// Python module implemented in Rust
#[pymodule]
fn rust_python_module(m: &mut PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(add, m)?)?;
    m.add_function(wrap_pyfunction!(fibonacci, m)?)?;
    Ok(())
}

Rebuild the module:

maturin develop

Test the Fibonacci function in Python:

import rust_python_module
import time

# Test the performance of the Rust implementation of the Fibonacci function
start_time = time.time()
result = rust_python_module.fibonacci(35)
end_time = time.time()

print(f"The 35th Fibonacci number is: {result}")
print(f"Rust implementation took: {end_time - start_time:.4f} seconds")

# Pure Python implementation of the Fibonacci function
def py_fibonacci(n):
    if n <= 0:
        return 0
    elif n == 1:
        return 1
    else:
        return py_fibonacci(n - 1) + py_fibonacci(n - 2)

# Test the performance of the Python implementation of the Fibonacci function
start_time = time.time()
result = py_fibonacci(35)
end_time = time.time()

print(f"Python implementation took: {end_time - start_time:.4f} seconds")

Run this script, and you will find that the Rust implementation of the Fibonacci function is significantly faster than the pure Python implementation!

Publishing Your Module

After completing module development, you can use maturin to publish it to PyPI:

# Build the wheel file
maturin build

# Upload to PyPI
maturin upload

Conclusion

Through this article, we learned how to build high-performance Python modules using Rust. This approach allows us to combine the ease of use of Python with the high performance of Rust, making it particularly suitable for projects that require handling compute-intensive tasks.

Using tools like PyO3 and maturin, we can easily bridge the two languages and fully leverage their respective advantages. Even if you are a beginner in Rust, this method is simple enough to try.

When you encounter performance bottlenecks in your next Python project, consider using Rust to solve them. This not only improves the execution efficiency of your program but also allows you to learn a new, vibrant programming language.

References

How to Build Python Modules in Rust: https://medium.com/@jdgb.projects/how-to-build-python-modules-in-rust-abb41bce60d0

Recommended Books

The second edition of “The Rust Programming Language” is an authoritative learning resource written by the Rust core development team and translated by members of the Chinese Rust community. It is suitable for all software developers who wish to evaluate, get started, improve, and research the Rust language, and is considered essential reading for Rust development work.

This book introduces the basic concepts of the Rust language to unique practical tools, covering advanced concepts such as ownership, traits, lifetimes, and safety guarantees, as well as practical tools like pattern matching, error handling, package management, functional features, and concurrency mechanisms. The book includes three complete project development case studies, guiding readers from scratch to develop practical Rust projects.

Notably, this book has been updated to include content from the Rust 2021 edition, meeting the systematic learning needs of beginners and serving as a reference guide for experienced developers, making it the best entry point for building solid Rust skills.