Writing Python Extensions in Rust with PyO3

Python is excellent for productivity but sometimes you need more speed. PyO3 lets you write Rust code that integrates seamlessly with Python. Here’s how.

Why Rust for Python Extensions

Rust gives you C-level performance with memory safety guarantees.

Setup

# Install maturin (build tool)
pip install maturin

# Create new project
maturin new my_rust_extension
cd my_rust_extension

Project structure:

my_rust_extension/
├── Cargo.toml
├── pyproject.toml
└── src/
    └── lib.rs

Hello PyO3

// src/lib.rs
use pyo3::prelude::*;

/// A function that greets a person
#[pyfunction]
fn greet(name: &str) -> PyResult<String> {
    Ok(format!("Hello, {}!", name))
}

/// A Python module implemented in Rust
#[pymodule]
fn my_rust_extension(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(greet, m)?)?;
    Ok(())
}

# Cargo.toml
[package]
name = "my_rust_extension"
version = "0.1.0"
edition = "2021"

[lib]
name = "my_rust_extension"
crate-type = ["cdylib"]

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

Build and test:

maturin develop
python -c "import my_rust_extension; print(my_rust_extension.greet('World'))"
# Hello, World!

Practical Example: Fast String Processing

Python:

def count_words(text: str) -> dict[str, int]:
    counts = {}
    for word in text.lower().split():
        word = word.strip('.,!?')
        counts[word] = counts.get(word, 0) + 1
    return counts

Rust:

use pyo3::prelude::*;
use pyo3::types::PyDict;
use std::collections::HashMap;

#[pyfunction]
fn count_words(py: Python, text: &str) -> PyResult<&PyDict> {
    let mut counts: HashMap<&str, usize> = HashMap::new();
    
    for word in text.to_lowercase().split_whitespace() {
        let clean = word.trim_matches(|c| matches!(c, '.' | ',' | '!' | '?'));
        *counts.entry(clean).or_insert(0) += 1;
    }
    
    let dict = PyDict::new(py);
    for (word, count) in counts {
        dict.set_item(word, count)?;
    }
    
    Ok(dict)
}

Classes in Rust

use pyo3::prelude::*;

#[pyclass]
struct Point {
    x: f64,
    y: f64,
}

#[pymethods]
impl Point {
    #[new]
    fn new(x: f64, y: f64) -> Self {
        Point { x, y }
    }
    
    fn distance(&self, other: &Point) -> f64 {
        let dx = self.x - other.x;
        let dy = self.y - other.y;
        (dx * dx + dy * dy).sqrt()
    }
    
    #[getter]
    fn x(&self) -> f64 {
        self.x
    }
    
    #[getter]
    fn y(&self) -> f64 {
        self.y
    }
    
    fn __repr__(&self) -> String {
        format!("Point({}, {})", self.x, self.y)
    }
}

#[pymodule]
fn geometry(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Point>()?;
    Ok(())
}

Usage:

from geometry import Point

p1 = Point(0.0, 0.0)
p2 = Point(3.0, 4.0)
print(p1.distance(p2))  # 5.0

Working with NumPy

use pyo3::prelude::*;
use numpy::{PyArray1, PyReadonlyArray1};

#[pyfunction]
fn sum_array(arr: PyReadonlyArray1<f64>) -> f64 {
    arr.as_slice().unwrap().iter().sum()
}

#[pyfunction]
fn multiply_array<'py>(
    py: Python<'py>, 
    arr: PyReadonlyArray1<f64>,
    factor: f64
) -> &'py PyArray1<f64> {
    let data: Vec<f64> = arr.as_slice()
        .unwrap()
        .iter()
        .map(|x| x * factor)
        .collect();
    PyArray1::from_vec(py, data)
}

Add to Cargo.toml:

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

Error Handling

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

#[pyfunction]
fn divide(a: f64, b: f64) -> PyResult<f64> {
    if b == 0.0 {
        Err(PyValueError::new_err("Cannot divide by zero"))
    } else {
        Ok(a / b)
    }
}

Async Support

use pyo3::prelude::*;
use pyo3_asyncio::tokio::future_into_py;
use tokio::time::{sleep, Duration};

#[pyfunction]
fn async_operation(py: Python) -> PyResult<&PyAny> {
    future_into_py(py, async move {
        sleep(Duration::from_secs(1)).await;
        Ok("Operation complete!")
    })
}

Performance Comparison

String processing benchmark:

import timeit
from my_rust_extension import count_words as rust_count

def python_count(text):
    counts = {}
    for word in text.lower().split():
        counts[word] = counts.get(word, 0) + 1
    return counts

text = open("large_file.txt").read()  # 10MB file

print(timeit.timeit(lambda: python_count(text), number=10))
# ~5 seconds

print(timeit.timeit(lambda: rust_count(text), number=10))
# ~0.3 seconds  (15x faster)

Parallelism with Rayon

use pyo3::prelude::*;
use rayon::prelude::*;

#[pyfunction]
fn parallel_sum(numbers: Vec<i64>) -> i64 {
    // Runs on all CPU cores
    numbers.par_iter().sum()
}

Python’s GIL is released during Rust execution, enabling true parallelism.

Build and Distribution

Development

maturin develop  # Install in current virtualenv

Wheels

maturin build --release  # Creates wheel file

Publishing

maturin publish  # Uploads to PyPI

Cross-compilation

# Build for multiple platforms
maturin build --release --target x86_64-unknown-linux-gnu
maturin build --release --target aarch64-apple-darwin

When to Use PyO3

Good For

• Compute-intensive loops
• String processing
• Numerical algorithms
• Parallelizable workloads
• Replacing C extensions

Probably Not Worth It

• Simple I/O operations
• One-off scripts
• When NumPy/Pandas already solve it
• Learning curve doesn’t pay off

Final Thoughts

PyO3 bridges Python’s productivity with Rust’s performance. The tooling (maturin) is excellent, and the integration is seamless.

For performance-critical Python code, it’s worth learning.

The best of both worlds: Python’s ease, Rust’s speed.