PyO3 による Rust の Python バインディング

こんにちは、 Gunosy Tech Lab 所属の ryoaita です。最近はスクラムマスターをしながら Rust を書いたりしてます。この記事は Gunosy Advent Calendar 2022 の 24 日目の記事です。

23 日目の記事は Liang さんの Gradle + Kotlin + CircleCIによるAndroid Google Playデプロイの自動化でした。

今回は Rust で Python モジュールを簡単に作成できる PyO3 を紹介します。

背景

現在、我々のプロジェクトではモブプロで Rust によるアプリケーションサーバーの開発を行っています。このサーバーは Python で学習した機械学習のモデルを Rust のコードで利用しています。この際、 Rust と Python で計算結果が一致することの保証を行う必要があります。これを簡単に実現するために、 PyO3 で Rust の Python バインディングを作成しました。

クイックスタート

PyO3 を導入するには maturin というコマンドを利用するのが簡単です。まずは maturin を pip でインストールします。

$ python -m pip install maturin

maturin new を実行すると Python のパッケージが生成されます。 hello-pyo3 というパッケージを作ってみます。

$ maturin new hello-pyo3

コマンドを実行すると、次のように質問されます。

Which kind of bindings to use? ›

ここで pyo3 を選択すると PyO3 を利用するプロジェクトが作成されます。

 ✨ Done! New project created hello-pyo3

コマンドが完了すると、 hello-pyo3 というディレクトリが作成され、次のようなファイルが配置されます。

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

1 directory, 3 files

それぞれのファイルの内容を見てみます。

[package]
name = "hello-pyo3"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "hello_pyo3"
crate-type = ["cdylib"]

[dependencies]
PyO3 = { version = "0.17.3", features = ["extension-module"] }

[build-system]
requires = ["maturin>=0.14,<0.15"]
build-backend = "maturin"

[project]
name = "hello-pyo3"
requires-python = ">=3.7"
classifiers = [
    "Programming Language :: Rust",
    "Programming Language :: Python :: Implementation :: CPython",
    "Programming Language :: Python :: Implementation :: PyPy",
]

use PyO3::prelude::*;

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

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

sum_as_string という Rust の関数を Python に提供する hello_pyo3 というパッケージが作成されました。試しに、 sum_as_string を実行してみましょう。 maturin develop を実行すると簡単に確認ができます。

maturin develop を利用するには virtualenv か conda を利用する必要があります。ここでは virtualenv を使ってみましょう。

$ maturin develop
💥 maturin failed
  Caused by: You need to be inside a virtualenv or conda environment to use develop (neither VIRTUAL_ENV nor CONDA_PREFIX are set). See https://virtualenv.pypa.io/en/latest/index.html on how to use virtualenv or use `maturin build` and `pip install <path/to/wheel>` instead.

venv という名前で virtualenv を作成します。

$ python -m venv venv
$ . venv/bin/activate

そして maturin develop を実行すると、wheel が作成され、venv の環境に hello_pyo3 がインストールされます。

試しに、 sum_as_string を実行してみると、次のような実行結果が得られます。

>>> import hello_pyo3
>>> hello_pyo3.sum_as_string(100, 1)
'101'

PyO3 が何をやっているか？

PyO3 は #[pyfunction] や #[pymodule] などマクロを使うと、Python に公開するためのグルーコードを生成します。

cargo-expand を使うと生成されるコードを見ることができます。実際にやってみましょう。

cargo install cargo-expand

cargo-expand の実行には nightly コンパイラが必要です。インストールしていない場合は rustup でツールチェインをインストールしてください。

rustup toolchain install nightly

cargo expand を実行するとマクロの展開結果が出力されます。

cargo expand > expanded.rs
    Checking hello-pyo3 v0.1.0 (/home/ryo/junks/hello-pyo3)
error: the option `Z` is only accepted on the nightly compiler
error: could not compile `hello-pyo3`

expanded.rs の内容は下のとおりです。

#![feature(prelude_import)]
#[prelude_import]
use std::prelude::rust_2021::*;
#[macro_use]
extern crate std;
use PyO3::prelude::*;
/// Formats the sum of two numbers as string.
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}
#[doc(hidden)]
mod sum_as_string {
    pub(crate) struct MakeDef;
    pub const DEF: ::PyO3::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
}
const _: () = {
    use ::PyO3 as _PyO3;
    impl sum_as_string::MakeDef {
        const DEF: ::PyO3::impl_::pyfunction::PyMethodDef = _PyO3::impl_::pymethods::PyMethodDef::fastcall_cfunction_with_keywords(
            "sum_as_string\0",
            _PyO3::impl_::pymethods::PyCFunctionFastWithKeywords(
                __pyfunction_sum_as_string,
            ),
            "Formats the sum of two numbers as string.\u{0}",
        );
    }
    unsafe extern "C" fn __pyfunction_sum_as_string(
        _slf: *mut _PyO3::ffi::PyObject,
        _args: *const *mut _PyO3::ffi::PyObject,
        _nargs: _PyO3::ffi::Py_ssize_t,
        _kwnames: *mut _PyO3::ffi::PyObject,
    ) -> *mut _PyO3::ffi::PyObject {
        let gil = _PyO3::GILPool::new();
        let _py = gil.python();
        _PyO3::callback::panic_result_into_callback_output(
            _py,
            ::std::panic::catch_unwind(move || -> _PyO3::PyResult<_> {
                const DESCRIPTION: _PyO3::impl_::extract_argument::FunctionDescription = _PyO3::impl_::extract_argument::FunctionDescription {
                    cls_name: ::std::option::Option::None,
                    func_name: "sum_as_string",
                    positional_parameter_names: &["a", "b"],
                    positional_only_parameters: 0usize,
                    required_positional_parameters: 2usize,
                    keyword_only_parameters: &[],
                };
                let mut output = [::std::option::Option::None; 2usize];
                let (_args, _kwargs) = DESCRIPTION
                    .extract_arguments_fastcall::<
                        _PyO3::impl_::extract_argument::NoVarargs,
                        _PyO3::impl_::extract_argument::NoVarkeywords,
                    >(_py, _args, _nargs, _kwnames, &mut output)?;
                let mut ret = sum_as_string(
                    _PyO3::impl_::extract_argument::extract_argument(
                        _PyO3::impl_::extract_argument::unwrap_required_argument(
                            output[0usize],
                        ),
                        &mut {
                            _PyO3::impl_::extract_argument::FunctionArgumentHolder::INIT
                        },
                        "a",
                    )?,
                    _PyO3::impl_::extract_argument::extract_argument(
                        _PyO3::impl_::extract_argument::unwrap_required_argument(
                            output[1usize],
                        ),
                        &mut {
                            _PyO3::impl_::extract_argument::FunctionArgumentHolder::INIT
                        },
                        "b",
                    )?,
                );
                if false {
                    use _PyO3::impl_::ghost::IntoPyResult;
                    ret.assert_into_py_result();
                }
                _PyO3::callback::convert(_py, ret)
            }),
        )
    }
};
/// A Python module implemented in Rust.
fn hello_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(
        {
            use sum_as_string as wrapped_pyfunction;
            ::PyO3::impl_::pyfunction::wrap_pyfunction(&wrapped_pyfunction::DEF, m)
        }?,
    )?;
    Ok(())
}
#[doc(hidden)]
mod hello_pyo3 {
    pub(crate) struct MakeDef;
    pub static DEF: ::PyO3::impl_::pymodule::ModuleDef = MakeDef::make_def();
    pub const NAME: &'static str = "hello_pyo3\u{0}";
    /// This autogenerated function is called by the python interpreter when importing
    /// the module.
    #[export_name = "PyInit_hello_pyo3"]
    pub unsafe extern "C" fn init() -> *mut ::PyO3::ffi::PyObject {
        DEF.module_init()
    }
}
const _: () = {
    use ::PyO3::impl_::pymodule as impl_;
    impl hello_pyo3::MakeDef {
        const fn make_def() -> impl_::ModuleDef {
            const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(
                hello_pyo3,
            );
            unsafe {
                impl_::ModuleDef::new(
                    hello_pyo3::NAME,
                    "A Python module implemented in Rust.\u{0}",
                    INITIALIZER,
                )
            }
        }
    }
};

Rust の強力なマクロのパワーでバインディングのためのコードが生成されました。

Rust のコードを Python に公開する

PyO3 による Python のモジュールの作成は#[pymodule] で行います。サブモジュールを作成した場合は、PyModule.add_submodule() を使います。

/// A Python module implemented in Rust.
#[pymodule]
fn hello_pyo3(_py: Python, m: &PyModule) -> PyResult<()> {
    register_child_module(py, m)?;
    Ok(())
}

fn register_child_module(py: Python<'_>, parent_module: &PyModule) -> PyResult<()> {
    let child_module = PyModule::new(py, "child_module")?;
    child_module.add_function(wrap_pyfunction!(func, child_module)?)?;
    parent_module.add_submodule(child_module)?;
    Ok(())
}

#[pyfunction]
fn func() -> String {
    "func".to_string()
}

>>> from parent_module import child_module
>>> child_module.function()
"func"

Rust の関数の Python への公開は #[pyfunction] で行います。

#[pyfunction]
fn greet(name: String) -> String {
    format!("hello {}!", name).to_string()
}

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

>>> from parent_module import greet
>>> greet("gunoguno")
"hello gunoguno!"

Rust で Python のクラスを作成することも可能です。 Python のクラスを作成するには #[pyclass] と #[pymethods] を利用します。

#[pyclass]
struct Accumulator {
    value: i32
}

#[pymethods]
impl Accumulator {
    #[new]
    fn new(value: i32) -> Self {
        Self { value }
    }

    fn increment(&mut self) -> i32 {
        self.value +=
    }
}

#[new] でアノテーションをしたメソッドはコンストラクタとして扱われます。

PyModule.add_class() でクラスをモジュールに追加できます。

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

Accumulator クラスが Python から使えるようになりました。

>>> import hello_pyo3
>>> hello_pyo3.Accumulator
<class 'builtins.Accumulator'>
>>> a = hello_pyo3.Accumulator(0)
>>> a.increment(1)
1
>>> a.increment(2)
3
>>> a.increment(3)
6

今後の課題

PyO3 を利用すると、簡単に Python のバインディングを作成できます。Rust のコードを変更すると、当然 Python のバインディングの修正が必要となります。そのため、修正が必要なコードが増えるため、保守のコストも増加します。そのため、PyO3 の導入に関する ADR では次のようなコンプライアンスを設定しています。