This post describes how I approached writing a Python extension in Rust. The post covers:
- why one would even want to do this 🙃
- the approaches for calling Rust code from Python
- an overview of how to create a Python module in Rust using PyO3
- some tricky parts, e.g. inheritance
- building and distributing wheels
Let’s get started.
There are two main reasons:
- To use Rust libraries that already exist, e.g. cryptography libraries.
- To do computationally intensive work that will be too slow in Python. Other approaches if this is the main motivation are using a C extension (e.g. as numpy does) or using projects like Cython or numba.
For my use case, I had the first reason, I wanted to prototype something using a Rust crate that implemented a cryptographic protocol.
There are multiple approaches for calling compiled Rust code from Python, including
PyO3. Here we’ll cover the two most popular:
cffi (considered easier to use than
You can use
extern keyword to allow other languages to call Rust functions.
For example this
add function is marked as a public external function using the
pub extern keywords. The
#[no_mangle] attribute just tells the compiler to preserve the readable function name.
Then, one can load this dynamic library and call the external Rust functions:
PyO3 is a very cool project that allows one to define a Python module entirely in Rust. The above example in PyO3 would be:
And to define the actual Python module:
This means that now from the Python interpreter we can just do:
As we can see, the PyO3 approach is very straightforward. You simply add attributes to structs and functions in Rust to indicate that they should be exposed to Python, and then you write a Rust function to indicate what the top-level functions and classes are for that module.
There are also similarly easy-to-use build tools (via
maturin) to handle the packaging and build process. You can see this example packaged using
setuptools-rust on GitHub here.
In the rest of this post, I’ll explain more about using PyO3.
The most important attributes to know are:
#[pyclass]: to expose a Rust struct as a Python class
#[pyfunction]: to expose a Rust function as a Python function
#[pymethods]: to expose the methods defined in an
implblock of a struct with the
#[pyclass]attribute as methods on the corresponding Python class
#[pymodule]: to expose a collection of structs or functions as a Python module
Using these attributes, PyO3 macros will do all the FFI work for you.
Functions and methods exposed to Python must have return values that are either native Rust types (that can be converted to
PyObject via the
ToPyObject trait) or Python object types (e.g.
dict). See the list of conversions here.
Functions that can fail should return
PyResult, which is a type alias for
Result<T, PyErr>. If the Err variant is returned, an exception will be raised on the Python side. Note that you can also create custom exception types.
Let’s create an example class, using
#[new] attribute is used for your object constructor and initialization logic in Python (equivalent of Python
In Python, you’d call
doris = Animal('Doris', 2, 0) to use this.
#[pyo3(get)] attribute lets one read
doris.name as member attributes. If you want to set attributes also, you can use
#[pyo3(get, set)] (which could replace the
Animal::feed() method if we wanted to).
We can add this class to a new module as follows:
The above parts can cover simple projects. Two more advanced topics we’ll cover are inheritance, and magic methods.
What if we want to make a subclasses, say, a
Lion, that inherits from
Animal? Here’s how we do it:
#[pyclass] annotations indicate the parent (
#[pyclass(subclass)]) and child (
#[pyclass(extends=Parent)]) classes. The tuple syntax in the return value of the child is a little “trick” intended for ergonomics: you return
PyResult<(Child, Parent)> or
(Child, Parent). PyO3 will then run
Into<PyClassInitializer> on the child, where PyClassInitializer is PyO3’s pyclass initializer.
One might be surprised to find that implementing magic methods doesn’t work in a
#[pymethods] impl block. It turns out that you can implement Python “magic” methods like
__richcmp__ using the
PyObjectProtocol trait and the
#[pyproto] attribute in a separate
impl block. For example, to add a nice string representation for
We want to build and distribute wheels that do not require the rust toolchain to be installed on target systems. Fortunately, with
maturin, that’s pretty simple. For
setup.py for the
zoo example would be:
See the full project here.
Locally, if we’re on macOS, to build macOS wheels:
python3 setup.py sdist bdist_wheel
To build manylinux wheels we can follow the procedure described in the setuptools-rust project. First we fetch the Python Packaging Authority manylinux image:
docker pull quay.io/pypa/manylinux2014_x86_64
Then using the default
build-wheels.sh script provided by
docker run --rm -v `pwd`:/io quay.io/pypa/manylinux2014_x86_64 /io/build-wheels.sh
This leaves us with built wheels in
dist/ ready for upload to PyPI. And we should just upload the manylinux wheels built by the script as PyPI does not support wheels with platform tags like
linux_x86_64 (these are also produced by the above wheel build command but can be discarded).
I hope you’re convinced that writing Rust extensions with PyO3 is approachable. To read more check out the PyO3 guide. If you want to see a larger example, you can check out the library I wrote using PyO3 here and install in Python 3.7+, via
pip install signal-protocol 😊 .