🦀 Oxidized Python 🐍

This repo is a small exploratory project for me to learn how Rust code can be used as a Python package with minimal overhead.

Setup instructions

These are my notes on the procedure to set up a python extension module usable for Python and Rust.

The core tools used to make this work are

PyO3

Rust bindings for Python

maturin

Building tool for rust binaries as python packages

Setup

Virtual environment

Everything needs to happen through an existing virtual environment. After creating this venv, make sure that everything you do is done with it activated. After creating the virtual environment we install the maturin package using pip.

# Specify name of the venv to use
venv_name="venv"

# Create a new python venv only if it does not already exist
if [ ! -d "$venv_name" ]
then
  python -m venv "$venv_name"
else
  echo "Virtual environment $venv_name already exists in $(pwd)"
fi

# Activate the newly created venv and install the `maturin` package
export VIRTUAL_ENV_DISABLE_PROMPT=1
. venv/bin/activate
python -m pip install --upgrade pip
pip install maturin

Project structure

Following the PyO3 examples, we can simply initialize our project in the current directory. Running maturin init is everything we need to do. This will set us up with all the necessary files and initial configuration.

When running the initialization script make sure to select pyo3 as the kind of bindings that we are going to use in this project. If you don’t specify the type of bindings using the --bindings flag maturin will ask you using an interactive prompt.

maturin init --bindings pyo3

✨ Done! Initialized project /Users/tbszlg/scm/rust/oxidized_python

After this initialization our current directory will be initialized with this project structure:

.
├── Cargo.toml
├── README.org
├── pyproject.toml
└── src
    └── lib.rs

As a nice bonus maturin automatically creates a github CI workflow in .github/workflows/CI.yml that will trigger the build of Python wheels for a bunch of target platforms and upload them as artifacts whenever a new tag is pushed to branches main or master.

PyO3 function and module definitions

The automatically created src/lib.rs file already holds some example code for how a Python binding to Rust code looks from the Rust side.

To create a function in Rust that should be usable through a Python interface we need to add the #[pyfunction] macro to the function. The automatically created example code already shows this for the sum_as_string function.

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

To add a function to a module importable by Python we use the similar #[pymodule] macro.

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

Using Rust implementations in Python

The easiest way to make the Rust implementations available to us in Python is by using the maturin develop command. Running it will build the package and install it in the existing virtual environment we created earlier.

The --quiet flag I’m using will disable all of the cargo build outputs. There is no harm in removing it.

maturin develop --quiet

🔗 Found pyo3 bindings
🐍 Found CPython 3.9 at /Users/tbszlg/scm/rust/oxidized_python/venv/bin/python
📡 Using build options features from pyproject.toml
💻 Using `MACOSX_DEPLOYMENT_TARGET=11.0` for aarch64-apple-darwin by default
📦 Built wheel for CPython 3.9 to /var/folders/w1/hjpns4zx3b75rrv40rpc40p80000gn/T/.tmpZ088ON/oxidized_python-0.1.0-cp39-cp39-macosx_11_0_arm64.whl
🛠 Installed oxidized_python-0.1.0

The outputs of this command tell us that maturin used pyo3 bindings and CPython 3.9 to build our project. It also built the wheel specifically for the target aarch64-apple-darwin, which is an important detail as each target platform requires its own build. Finally we see that it built a wheel file oxidized_python-0.1.0-cp39-cp39-macosx_11_0_arm64.whl and installed it into our current python venv.

To test things we can run a tiny python snippet to validate that the module and function implemented in Rust are now available to us in Python. We should be able to import the oxidized_python module and call the sum_as_string function defined there.

import oxidized_python


res = oxidized_python.sum_as_string(3, 7)
print(f"The output of `sum_as_string` is {res} and of type {type(res)}")

Perfect! Our Python interpreter was able to import the functions defined in our Rust code, call them and work with the results as expected!

Dual Rust & Python library

The way we set up the project until now allows us to create a Python library from Rust code. But what if we would also want to make the library available to other Rust implementations?

In that case we simply add the specification to the Cargo.toml file that instructs the build process to additionally compile the package into a Rust library.

Adding the "lib" specification below is enough to do so.

[lib]
name = "oxidized_python"
crate-type = ["cdylib", "lib"]

Face recognition package

Preparations

• Download model files from https://github.com/deepinsight/insightface/tree/master/python-package#model-zoo
• Unzip into assets/models/ directory (onnx files)

Testing Python integration

This small snippet is enough to run the face-recognition implementation in Rust via the Python interpreter.

import cv2
import numpy as np

import oxidized_python


tom = cv2.imread("assets/tom.png", cv2.IMREAD_COLOR).astype(np.uint8)
tom = cv2.cvtColor(tom, cv2.COLOR_BGR2RGB)
hanks = cv2.imread("assets/hanks.jpg", cv2.IMREAD_COLOR).astype(np.uint8)
hanks = cv2.cvtColor(hanks, cv2.COLOR_BGR2RGB)

environment = oxidized_python.PyAgentEnvironment()
agent = environment.create_agent()

similarity = agent.get_similarity(tom.tolist(), hanks.tolist())

print(similarity)

0.7528126835823059

Testing Rust usage

The developed library can also be used by Rust. For this purpose I implemented a binary in src/main.rs that instantiated a face-recognition agent and computes the similarity between the two images.

To compile the binary it is enough to run the regular cargo build command:

cargo build --bin oxidized_python

The binary is implemented such that it accepts two command line arguments - the path to two images to calculate face similarity for.

./target/debug/oxidized_python ./assets/tom.png ./assets/hanks.jpg

Similarity: 0.75250286

As we can see the outcome similarity is the same between the usage of the python wrapper or the direct Rust implementation.

Latency comparisons