Include Rust in your Python Project using Poetry and Maturin
Let’s try and include some Rust code in a Python project to see how it’s done.
Warning
I attempted this because I was concerned while building games and simulations in Python. I was wondering if I’d hit a point in which Python might not be performant for a given problem and I’d need to port some functionality to another language. I didn’t want to get too far into it and run into problems, so I pieced together a trivial example of something I might come across just to see what it’s like to include Rust in Python.
- It’s extremely easy to include Rust code into your Python project.
- Don’t use this for everything.
- This post is going to show a trivial example that could improve something you’re working on. Is it useful code? Probably not.
High Level Steps
- Install python dependencies
- Build a Rust crate that can be turned into a Python wheel using Maturin
- Include the wheel into the package with Poetry
- Do a rough comparison in performance of Python vs. Rust function
Project Structure Before
We’re going to have a Python virtual environment, with dependencies including: poetry and maturin. My project is going to be called “test_app”. My Rust crate will be called “collision”.
This is the structure before we introduce Rust. Note: I don’t know if there’s any conventional way to structure your Python app with Rust code, this is just how I chose to do it to keep things organized.
test_app
src
test_app
__init__.py
main.py
rust_crates
collision
pyproject.toml
README.md
Code
Create the structure and get poetry installed so you can install other dependencies.
mkdir -p test_app/src/test_app && cd test_app
mkdir -p src/rust_crates/collision
# Make sure the necessary files exist
touch README.md
touch pyproject.toml
touch src/test_app/__init__.py
touch src/test_app/main.py
# Set up Python
python3.11 -m venv venv && source venv/bin/activate
pip install poetry
pyproject.toml
Add the following into pyproject.toml to set up the necessary dependencies.
[tool.poetry]
name = "test_app"
version = "0.1.0"
description = "Test App"
authors = ["First Last <email@example.com>"]
readme = "README.md"
packages = [{ include = "test_app", from = "src/" }]
[tool.poetry.dependencies]
python = ">=3.11"
# collision = { file = "src/rust_crates/collision/target/wheels/collision-0.1.0-cp311-cp311-macosx_11_0_arm64.whl" }
maturin = "^1.4.0"
[build-system]
requires = ["poetry>=1.0"]
build-backend = "poetry.masonry.api"
Above is the necessary configuration to get started. After building our Rust code, be sure to come back and uncomment the collision dependency.
Install the dependencies by running:
poetry lock && poetry install
src/test_app/main.py
Build a simple Python script that iterates over a collision detection function. It’s definitely not useful, but illustrates a cpu-bound operation with the loop. We will expand on this later with the Rust code.
import time
def benchmark_collision_detection(x1, y1, r1, x2, y2, r2, iterations=10000):
# If the square of the distance is less than the square of the sum of the radii
# it means the circles overlap, indicating a collision.
for _ in range(iterations):
_ = ((x1 - x2) ** 2 + (y1 - y2) ** 2) < (r1 + r2) ** 2
def main():
iterations = 10000
start_time = time.time()
python_time = benchmark_collision_detection(100, 100, 30, 150, 150, 40, iterations)
python_time = time.time() - start_time
print(f"{iterations} iterations in Python took {python_time:.6f} seconds")
if __name__ == "__main__":
main()
Make sure it works with: python -m test_app.main
Configure a Rust Crate with Maturin
cd src/rust_crates/collision
maturin init # hit enter to select pyo3
You should not have additional files created with some example Rust code in src/lib.rs. Let’s modify this to do something similar to what our Python code did.
src/lib.rs
Here we add a benchmark_collision_detection function similar to how we implemented it in Python. We expose this Rust function to Python with a combination of attributes and macros explained later.
use pyo3::prelude::*;
use std::time::Instant;
fn check_collision(x1: f32, y1: f32, r1: f32, x2: f32, y2: f32, r2: f32) -> bool {
let distance_squared = (x1 - x2).powi(2) + (y1 - y2).powi(2);
distance_squared < (r1 + r2).powi(2)
}
#[pyfunction]
fn benchmark_collision_detection(
x1: f32,
y1: f32,
r1: f32,
x2: f32,
y2: f32,
r2: f32,
iterations: usize,
) -> PyResult<f64> {
let start = Instant::now();
for _ in 0..iterations {
check_collision(x1, y1, r1, x2, y2, r2);
}
let duration = start.elapsed().as_secs_f64();
Ok(duration)
}
#[pymodule]
fn collision(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(benchmark_collision_detection, m)?)?;
Ok(())
}
#[pyfunction]: This attribute marks the following function as a Python callable. This means the function benchmark_collision_detection can be called from Python.#[pymodule]: This attribute indicates that the following function is used for initializing a Python module.m.add_function(wrap_pyfunction!(benchmark_collision_detection, m)?)?;: This line adds thebenchmark_collision_detectionfunction to the module. Thewrap_pyfunction!macro is used to generate a wrapper that allows the function to be called from Python.
Build a Python wheel from the Rust crate
From the root of the collision directory, build this Rust crate into a Python wheel with the following:
maturin build # creates target/wheels/collision-0.1.0-*.whl
To increase the version of this wheel, update Cargo.toml with a new version before running the build.
Include Rust in Python
Below is what your project structure should look like now that you’ve built a Python wheel.
test_app
src
test_app
__init__.py
main.py
rust_crates
collision
src
lib.rs
target
wheels
*.whl # e.g. collision-0.1.0-cp311-cp311-macosx_11_0_arm64.whl
Cargo.toml
pyproject.toml
pyproject.toml
README.md
Go back into the <root>/pyproject.toml and uncomment the collision dependency, making sure to change the path to the wheel that was generated for you. This will differ based on the system you’re building from.
src/test_app/main.py
Update main.py to call the Rust code
import time
import collision # rust code
def benchmark_collision_detection(x1, y1, r1, x2, y2, r2, iterations=10000):
# If the square of the distance is less than the square of the sum of the radii
# it means the circles overlap, indicating a collision.
for _ in range(iterations):
_ = ((x1 - x2) ** 2 + (y1 - y2) ** 2) < (r1 + r2) ** 2
def main():
iterations = 10000000
# python
start_time = time.time()
python_time = benchmark_collision_detection(100, 100, 30, 150, 150, 40, iterations)
python_time = time.time() - start_time
print(f"{iterations} iterations in Python took {python_time:.6f} seconds")
# rust
start_time = time.time()
python_time = collision.benchmark_collision_detection(100, 100, 30, 150, 150, 40, iterations)
rust_time = time.time() - start_time
print(f"{iterations} iterations in Rust took {rust_time:.6f} seconds")
if __name__ == "__main__":
main()
Install and run the app.
poetry lock && poetry install
python -m test_app.main
Experiment with the number of iterations and what each function is doing and you will likely see that Rust is more performant than Python for cpu-bound operations.
Takeaways
Keep the following in mind, otherwise adding Rust to your Python might just make it slower.
- Port entire loops to Rust, don’t just iterate in Python and hope single function calls will be faster in Rust.
- Similarly, don’t generate a list of objects in Python only to pass them all into Rust. Generate variables natively in Rust if you can.
- Benchmark your code first to find where it’s running hot. Try and port that code first, leaving the rest in Python.