How to write Python extensions in Rust with PyO3

Py03 lets you combine Rust's speed and memory safety with Python's ease of use. Get started writing Rust extensions for Python that work just like regular Python modules.

tadamichi/Shutterstock
By Serdar Yegulalp
Senior Writer, InfoWorld
Feb 22, 2023 2:00 AM

Every programming language has strengths and weaknesses. Python offers many convenient programming conventions but is computationally slow. Rust gives you machine-level speed and strong memory safety but is more complex than Python. The good news is, you can combine the two languages, wielding Python's ease of use to harness Rust's speed and power. The PyO3 project lets you leverage the best of both worlds by writing Python extensions in Rust.

With PyO3, you write Rust code, indicate how it interfaces with Python, then compile Rust and deploy it directly into a Python virtual environment, where you can use it unobtrusively with your Python code.

This article is a quick tour of how PyO3 works. You'll learn how to set up a Python project with a PyO3 create, how to expose Rust functions as a Python module, and how to create Python objects like classes and exceptions in Rust.

Setting up a Python project with PyO3

To start creating a PyO3 project, you need to begin with a Python virtual environment, or venv. This is not just for the sake of having your Python project organized, but also to provide a place to install the Rust crate you'll be building with PyO3. (If you haven't already installed the Rust toolchain, do that now.)

The exact organization of the project directories can vary. In the examples shown in PyO3's documentation, the PyO3 project is built in a directory that contains the Python project and its virtual environment. Another approach is to create two subdirectories: one for your Python project and its venv, and the other for the PyO3 project. The latter approach makes it easier to keep things organized, so we'll do that:

  1. Create a new directory to hold both your Python and Rust projects. We'll call them pyexample and rustexample, respectively.
  2. In the pyexample directory, create your virtual environment and activate it. We'll eventually add some Python code here. It's important that you perform all your work with both the Rust and Python code in your activated venv.
  3. In your activated venv, install the maturin package with pip install maturin. maturin is the tool we use to build our Rust project and integrate it with our Python project.
  4. Switch to the Rust project directory and type maturin init. When asked what bindings to select, choose pyo3.
  5. maturin will then generate a Rust project in that directory, complete with a Cargo.toml file that describes the project. Note that the project will be given the same name as the directory it's placed in; in this case it'll be rustexample.

Rust functions in a PyO3 project

When you create a PyO3 project's scaffolding with maturin, it auto-creates a code stub file in src/lib.rs. This stub contains code for two functions—a single sample function, sum_as_string, and a function named after your project that exposes other functions as a Python module.

Here's an example sum_as_string function:

#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

The #[pyfunction] macro, from the pyo3 crate, indicates a given function is to be wrapped with an interface to Python. The arguments it takes in and the results it returns are all translated from and to Python types automatically. (It's also possible to specify Python-native types to take in and return; more on this later.)

In this example, sum_as_string takes in two arguments that must be translatable to a Rust-native 64-bit integer. For such a case, a Python program would pass in two Python int types. But even then, you'd have to be careful: those int types would need to be expressable as a 64-bit integer. If you passed 2**65 to this function, you'd get a runtime error because a number that big can't be expressed as a 64-bit integer. (We'll talk about another way to get around this limitation later.)

The return value for this function is a native Python type—a PyResult object that contains a String. The last line of the function returns a String, which the PyO3 wrapper automatically wraps as a Python object.

It's also possible for pyfunction to describe the signature that a given function will accept—for instance, if you want to accept multiple positional or keyword arguments.

Python and Rust types in PyO3 functions

You'll want to get famliar with how Python and Rust types map to each other, and make some choices about what types to use.

Your function can accept Rust types that are converted automatically from Python types, but this means containers like dictionaries must be converted entirely at the function boundary. That might be slow if you pass a large object, such as a list with thousands of objects. To that end, this is best done if you're passing a single value, like an integer or a float, or container objects you know aren't going to have many elements.

You can also accept Python-native types at the function boundary, and use Python-native methods to access them within the function. This is faster at the function boundary, so it's a better choice if you're passing container objects with an indeterminate number of elements. But accessing container objects requires using Python-native methods that are bound by the GIL (Global Interpreter Lock), so you'll need to convert any values from the object into Rust-native types for speed.

Python modules in a PyO3 project

pyfunction functions by themselves aren't directly exposed to Python by way of a module. To do this, we need to create a Python module object through PyO3 and expose our pyfunction functions through it.

The lib.rs file already has a basic version created for you, which looks like this:


#[pymodule]
fn rustexample(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

The pymodule macro indicates the function in question will be exposed as a module to Python, with the same name (rustexample). We take each of the previously defined functions and expose them through the module using the .add_function method. This may seem a little boilerplate, but it provides flexibility when creating the module—for example, by allowing you to create submodules if needed.

Compiling a PyO3 project

Compiling your PyO3 project for use in Python is generally pretty simple:

  1. If you haven't done so already, activate the virtual environment where you installed maturin.
  2. Set your Rust project as your current working directory.
  3. Run the command maturin dev to build your project.

The results should look something like this:


(.env) PS D:\Dev\pyo3-article\rustexample> maturin dev -r
    Updating crates.io index
    [ ... snip ... ]
  Downloaded 10 crates (3.2 MB) in 2.50s (largest was `windows-sys` at 2.6 MB)
🔗 Found pyo3 bindings
🐍 Found CPython 3.11 at D:\Dev\pyo3-article\pyexample\.env\Scripts\python.exe
   [ ... snip ... ]
   Compiling rustexample v0.1.0 (D:\Dev\pyo3-article\rustexample)
    Finished release [optimized] target(s) in 10.86s
📦 Built wheel for CPython 3.11 to [ ... snip ...]
    \.tmpUbXtlF\rustexample-0.1.0-cp311-none-win_amd64.whl
🛠 Installed rustexample-0.1.0

By default, maturin builds Rust code in pre-release mode. In this example, we passed the -r flag to maturin to build Rust in release mode.

The resulting code should then be installed directly in your virtual environment, and you should be able to see it with pip list:


(.env) PS D:\Dev\pyo3-article\rustexample> pip list
Package     Version
----------- -------
maturin     0.14.12
pip         23.0
rustexample 0.1.0
setuptools  67.1.0

To test out your built package, launch the Python instance in your virtual environment and try importing the package:


Python 3.11.1 (tags/v3.11.1:a7a450f, Dec  6 2022, 19:58:39)
 [MSC v.1934 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import rustexample
>>> rustexample
<module 'rustexample' from 'D:\\Dev\\pyo3-article\\pyexample\\
    .env\\Lib\\site-packages\\rustexample\\__init__.py'>

It ought to import and run like any other Python package.

Advanced PyO3

So far, you've seen only the very basics of what PyO3 can do. But PyO3 supports a good many other Python features, many of which you will likely want to interface with Rust code.

Big integer support

Python automatically converts integers to "big integers," or integers of arbitrary size. If you want to pass a Python integer object into a PyO3 function and use it as a Rust-native big integer, you can do this with pyo3::num_bigint, which uses the existing num_bigint crate. Just remember that big integers might not support all operations.

Parallelism

As with Cython, any purely Rust code that doesn't touch the Python runtime can be run outside of the Python GIL. You can wrap such a function in the Python::allow_threads method to suspend the GIL while it executes. Again, this has to be purely Rust code with no Python objects in use.

Holding the GIL with Rust lifetimes

PyO3 provides a way to hold the GIL by way of Rust's lifetimes mechanism, which gives you a way to take either mutable or shared access to Python objects. Different object types have different GIL rules.

You can access a generic Python object with the PyAny type, or you can use more precise types like PyTuple or PyList. These are a little faster, since PyO3 can generate code specific to those types. No matter which types you use, you should assume you need to hold the GIL for the entire time you're working with the object.

If you want a reference to a Python object outside the GIL—for instance, if you're storing a Python object reference in a Rust struct—you can use the Py<T> or PyObject (essentially Py<PyAny>) types.

For a Rust object wrapped in a (GIL-holding) Python object—yes, this is possible!—you can use PyCell<T>. You'd typically do this if you wanted to access the Rust object while maintaining its Rust aliasing and reference rules. In that case, the wrapping Python object's behavior doesn't interfere with what you want to do. Likewise, you can use PyRef<T> and PyRefMut<T> to get borrowing references, static and mutable, to such objects.

Classes

You can define Python classes in PyO3 modules. If you add the #[pyclass] attribute to a Rust struct or a fieldless enum, they can be treated as the basic data structure for a class. To add instance methods, you'd use #[pymethods] with an impl block for the class that contains the functions to use as methods. It's also possible to create class methods, attributes, magic methods, slots, callable classes, and many other common behaviors.

Keep it in mind that Rust's behaviors impose some limitations. You can't provide lifetime parameters for classes; they all have to work as 'static. You also can't use generic parameters on types being used as Python classes.

Exceptions

Python exceptions in PyO3 can be created in Rust code with the create_exception! macro, or by importing one of a few predefined standard exceptions with the import_exception! macro. Note that, as with functions, you have to manually add PyO3-created exceptions to a module to make them available to Python.

Conclusion

For a long time, building Python extensions typically meant learning C with all its minimalism and lack of native safeties. Or, you could use a tool like Cython with all its idiosyncrasies. But for developers who already know Rust and want to use it hand-in-hand with Python, PyO3 provides a convenient and powerful way to do it.