PyO3でPythonからRustを使う

PyO3を試した
https://github.com/PyO3/pyo3
https://github.com/PyO3/setuptools-rust
https://pyo3.rs/v0.14.5/

version
Python: 3.8.10
Rust: 1.55.0

OS は Windows

ライブラリを作る

cargo new --lib addarray

Carto.toml に [lib] と [dependencies.pyo3] を追加

[package]
name = "addarray"
version = "0.1.0"
edition = "2018"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]

[lib]
name = "add_array_rs_test"
crate-type = ["cdylib"]

[dependencies.pyo3]
version = "0.14.5"
features = ["extension-module"]

src/lib.rs の中身

use pyo3::prelude::*;

#[pyfunction]
fn add_array_rs(n: u64, x: u64) -> PyResult<u64> {
    let mut a = vec![0u64; n as usize];
    for _ in 0..x {
        for i in 0..n as usize {
            a[i] += 1;
        }
    }
    Ok(a.iter().sum())
}

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

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn it_works() {
        assert_eq!(add_array(100, 100), 10000);
    }
}

setup.py, pyproject.toml, MANIFEST.in を作成

setup.py

from setuptools import setup
from setuptools_rust import Binding, RustExtension

setup(
    name="hello-rust",
    version="1.0",
    rust_extensions=[RustExtension(
        "add_array_rs_test", binding=Binding.PyO3)],
    zip_safe=False,
)

pyproject.toml

[build-system]
requires = ["setuptools", "wheel", "setuptools-rust"]

MANIFEST.in

include Cargo.toml
recursive-include src *

以下のように setup すると

python ./setup.py develop

以下のようになって

Installed (省略)\addarray
Processing dependencies for hello-rust==1.0
Finished processing dependencies for hello-rust==1.0

add_array_rs_test.cp38-win_amd64.pyd ができる

この .pyd を適切な場所に置けばよい

python 版と比較

import time
from add_array_rs_test import add_array_rs


def add_array(n: int, x: int) -> int:
    a = [0] * n
    for _ in range(x):
        for i in range(n):
            a[i] += 1
    return sum(a)


def add_array_2(n: int, x: int) -> int:
    a = [0] * n
    for _ in range(x):
        a = [i + 1 for i in a]
    return sum(a)


t1 = time.time()
assert add_array(10000, 10000) == 100000000
t2 = time.time()
print(f"py:  {t2 - t1:.3f} sec")

t1 = time.time()
assert add_array_2(10000, 10000) == 100000000
t2 = time.time()
print(f"py2: {t2 - t1:.3f} sec")

t1 = time.time()
assert add_array_rs(10000, 10000) == 100000000
t2 = time.time()
print(f"rs:  {t2 - t1:.3f} sec")

結果（1回しか計ってないけど）

py:  7.786 sec
py2: 4.403 sec
rs:  5.008 sec

最適化してないので遅い。

python ./setup.py install

とすると

Installed (省略)\python38\lib\site-packages\hello_rust-1.0-py3.8-win-amd64.egg
Processing dependencies for hello-rust==1.0
Finished processing dependencies for hello-rust==1.0

となって

py:  8.398 sec
py2: 4.448 sec
rs:  0.047 sec

爆速

おっと、setup いらなかった？

cargo build --release

target/release/add_array_rs_test.dll ができているので、拡張子を .pyd に変えて適切な場所に置けばよい。

これを書く前に試したときは dynamic module does not define module export function となったのだけど。

関連記事

• PyO3でPythonからRustを使う(1.1)
• PyO3でPythonからRustを使う(2: 並列処理)

addarray の元ネタ