Making Your Python Code Faster Using Rust

Python is an interpreted and dynamically typed programming language, that has become increasingly popular for many reasons, a few of which include its simplicity compared to other programming languages such as C, C++, and Java. It allows you to accomplish more with less code due to its wide range of libraries.

Additionally, it has a very large and supportive community, making it a preferred choice for academic and research work.

While these unique features of Python allow for rapid application development, they also come with performance costs. 

Python is slow.

Here is a quick illustration that shows how slow Python can be: Below is a Python function that calculates the sum of numbers from 1 to any specified number — in our case, we used 10 million.

import time

def sum_numbers(n):

    total = 0

    for i in range(1, n+1):

        total += i

    return total

start_time = time.time()

result = sum_numbers(10**7)

end_time = time.time()

print("Python: Sum =", result)

print("Time taken:", (end_time - start_time) * 1000, "milliseconds")

And here is also the Rust code to perform the same task:

use std::time::Instant;

fn sum_numbers(n: i64) -> i64 {

    let mut total = 0;

    for i in 1..=n {

        total += i;

    }

    total

}

fn main() {

    let start_time = Instant::now();

    let result = sum_numbers(10_000_000);

    let end_time = Instant::now();

    println!("Rust: Sum = {}", result);

    println!("Time taken: {:.2} milliseconds", (end_time - start_time).as_secs_f64() * 1000.0);

}

Here is the result when I ran both of them on the same machine:

While this can’t be considered a perfect benchmark measurement, it gives you an idea. From the screenshot above, you’ll notice that there is a huge difference between the time it took to complete the execution of the Python and the Rust code, over 400% 🙂 

So, in the rest of this guide, we’ll see how you can speed up your Python code with Rust by writing the code in Rust and using it as a Python module in your Python project. 

Rust is a memory-safe and highly performant programming language. It’s loved among programmers because of how it focuses on performance and ensures you write robust software that is less likely to be buggy. By writing code in Rust and using it in Python, we can benefit from the performance that Rust brings to the table.

We’ll start by setting up a Rust project and then we’ll test out our code in Python and we’ll see the difference in performance. Come on, let’s get into it.

Setting up the Rust project

To follow along with this guide, you’ll need to install Python from Python.org and Rust with Rustup

Once that is done, you’ll have cargo, the Rust package manager installed as well. We’ll set up our Rust project with cargo by running the command below:

cargo init --lib

The code above will initialize a Rust library application setup with the following directory structure:

.
├── Cargo.toml
└── src
    └── lib.rs

So, we’ll write our custom code in the lib.rs file, you can delete the default content in it and get ready for when we’ll write that code  

Next, we’ll install the required dependencies that will allow us to implement Python modules in our Rust code and add a little bit of configuration to make that work.

Introducing maturin

Technically, all we need is maturin and pyo3. Maturin is an open-source Python project that allows you to build and publish Rust crates with pyo3, while PyO3 enables you to generate native Python modules. So, pyo3 will generate is like the library to generate the modules, while maturin is like the package manager that does all the job of building and publishing it.

There are several ways to install maturin, you could use Python Pip  since you already have Python installed by running the code below:

python3 -m venv venv  && source ./venv/bin/activate && pip install maturin

The command above creates a virtual environment, activates it, and installs maturin in that environment.

You can also use other methods to install maturing as more options are outlined in the documentation.

Configuring Cargo.toml

We’ll need to configure the project, add the dependencies, and define the library as a C dynamic cdylib, when you set the crate type to cdylib, you’re telling Cargo to build the Rust code as a dynamic library that can be linked with other languages, such as C or Python.

So, in our Cargo.toml file, add the code below:

[lib]

name = "sum_numbers"

crate-type = ["cdylib"]

path = "src/lib.rs"

The name sum_numbers in the [lib] configuration is how this module will be referred to in your Python code. While the crate-type [“cdylib”] is necessary to produce a shared library for Python to import from and finally, the path tells us where to find the library code.

Next, we’ll add the dependencies. For this project, we only have one dependency, for more complex projects you might have more.

[dependencies]

pyo3 = {version = "0.18.0", features= ["abi3-py37", "extension-module"]}

This is where we add the pyo3 crate and enable some necessary features like the “abi3-py37” which indicates that the extension module being built is compatible with Python 3.7’s ABI (Application Binary Interface) and the extension-module feature enables the necessary features needed for pyo3 to build an extension module.

To conclude our configuration, we’ll add a [build-system] to the Cargo.toml file:

[build-system]
requires = ["maturin"]
build-backend = "maturin"

Which will require maturin and set maturin as the backend for this library. In other words, when you compile your project, it will use maturin as the tool to manage the Rust code and create the Python extension module.

That’s all there is about the configuration.

Writing the Rust code

Now that we have the configuration all setup, let’s get on with writing the Rust code that will boost our Python app and actually run it. 

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

#[pyfunction]
fn sum_numbers(n: i64) -> i64 {
    let mut total = 0;
    for i in 1..=n {
        total += i;
    }
    total
}
#[pymodule]
fn get_sum(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_numbers, m)?)?;
    
    Ok(())
}

The code above is similar to our previous Rust code, the only obvious difference is that we added pymodule and pyfunction attributes to the function and then we added PyResult<()> as the return type.

Let’s describe what those parameters represent:

First off, the #[pyfunction] attribute is a Rust attribute from the pyo3 library that allows you to mark a Rust function to be used in Python. The #[pymodule] defines a Python module and lastly

this line  m.add_function(wrap_pyfunction!(sum_numbers, m)?)?; adds the Rust function to the Python module get_sum.

The rest of the code is a normal Rust code. This is the fundamentals of creating a Python module and using it in your Python code.

Using the Rust code in Python

Now, let’s get to using it in your Python project. To do that, we’ll need to create a development build of the library by running the command below:

maturin develop

Make sure you are in the virtual environment before running the code above, if you installed maturin in your virtual environment. The output should look like this:

Once that is done, the get_sum module will now be installed in your Python virtual environment and we can now go ahead to use it. Create a perf.py file and write the code to use the library, add the code below into the perf.py file:

import time
from get_sum import sum_numbers

start_time = time.time()
result = sum_numbers(10_000_000)
end_time = time.time()

print("Rust: Sum =", result)
print("Time taken: {:.2f} milliseconds".format((end_time - start_time) * 1000.0))

The code is simple, we just import the module and pass the number we intend to create it’s summation. Now, we can run it with python perf.py, you’ll see a very clear difference between the result we’ll get with this and the result we got earlier.

We just increased our code performance by over 350%. That for me is an amazing feat. We’ve succeeded in seeing up our Python code with Rust.

Our module was only installed in our Python virtual environment when we ran maturin develop and we can’t distribute it to our team members. Let’s go ahead and create a build for distribution by running maturin build

Notice the last item in the result with the extension .wl. We can now go ahead to distribute the file ./target/wheels/app-0.1.0-cp38-abi3-macosx_10_7_x86_64.whl  for installation. You can even try installing it yourself by running pip install ./target/wheels/app-0.1.0-cp38-abi3-macosx_10_7_x86_64.whl 

Since I already have it installed, it will just say, I’ve already installed it:

You can also publish it to Pypi by running maturin publish

That’s all there is to it

Conclusion

Python is a powerful programming language, but it has shortcomings like every other language. It feels so good to be able to leverage the speed of Rust to make Python even faster. This guide has outlined the basic process to leverage Rust to improve Python performance; however, we have just scratched the surface. There is more— you should read the maturin and pyo3 documentation to explore more complex scenarios.

Happy Hacking!