PyO3でRust製アプリケーションにPythonを組み込む

製品開発グループの大久保です。aptpod Advent Calendar 2024の12月5日の記事を担当します。

社内では、エッジ側でintdashに接続可能なデバイスを簡単に開発するためデバイス開発キットの基盤製品を開発しており、それをDevice Connectorと呼んでいます。これは主にRustによって開発しています。

Rustは高速なコンパイル型言語であり、デバイスとの通信、制御に適しています。しかしながら、データのフィルタリングなどの細かい挙動を、再ビルドを必要とせず利用者が多くて学習コストの小さいPythonで記述したいというニーズもあります。そのような、アプリケーションの基本部分はRustで実装しながら、状況に応じて制御を変えたい部分はPythonで記述したいという場合、RustからPythonを呼び出すことになります。それを今回は PyO3で試してみたいと思います。

PyO3とは

PyO3 は、PythonをRustを利用するためのバインディングを提供します。RustからPythonを呼び出すだけでなく、Pythonから利用されるモジュールをRustで記述することができます。

大体の使い方は以下のユーザーガイドを見れば分かると思います。

https://pyo3.rs/main/getting-started.html

RustからPythonを呼び出す

実際にRustからPythonを呼び出してみましょう。実行時には libpython3.x.so の機能を呼び出すので、例えばUbuntuなら以下のように依存ライブラリをインストールする必要があります。

sudo apt install python3-dev

Cargo.tomlに以下の依存関係を記述します。

[dependencies]
pyo3 = { version = "0.22.6", features = ["gil-refs"] }

main.rs は以下のように記述します。

#[pyclass(name = "Greeting")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct PyGreeting {
    text: String,
}

#[pymethods]
impl PyGreeting {
    #[getter]
    fn get_text(&self) -> &str {
        &self.text
    }
}

#[pyclass(name = "Response")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct PyResponse {
    text: String,
}

#[pymethods]
impl PyResponse {
    #[new]
    fn new(text: String) -> Self {
        Self { text }
    }
}

#[pymodule]
fn mymodule(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<PyGreeting>()?;
    module.add_class::<PyResponse>()?;
    Ok(())
}

fn main() {
    pyo3::append_to_inittab!(mymodule);
    pyo3::prepare_freethreaded_python();

    let script = r#"
import mymodule

def greeting(g):
    if g.text == "hello":
        res = "こんにちは"
    else:
        res = "なんですか?"
        
    return mymodule.Response(res)
"#;

    println!(
        "{:?}",
        run_script(
            script,
            PyGreeting {
                text: "hello".into()
            }
        )
    );
    println!(
        "{:?}",
        run_script(
            script,
            PyGreeting {
                text: "hoge".into()
            }
        )
    );
}

fn run_script(script: &str, greeting: PyGreeting) -> PyResponse {
    let result: PyResult<PyResponse> = Python::with_gil(|py| {
        let pyfunc: Py<PyAny> = PyModule::from_code_bound(py, &script, "", "")?
            .getattr("greeting")?
            .into();
        let result = pyfunc.call1(py, (greeting,));
        let response = match result {
            Ok(response) => response,
            Err(e) => {
                eprintln!("{}", e);
                if let Some(e) = e.traceback_bound(py) {
                    if let Ok(e) = e.into_gil_ref().format() {
                        eprintln!("{}", e);
                    }
                }
                panic!();
            }
        };
        let response: PyResponse = response.extract(py)?;
        Ok(response)
    });
    result.unwrap()
}

以下細かく解説します。

#[pyclass(name = "Greeting")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct PyGreeting {
    text: String,
}

#[pymethods]
impl PyGreeting {
    #[getter] // getterとして定義
    fn get_text(&self) -> &str {
        &self.text
    }
}

#[pyclass(name = "Response")]
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct PyResponse {
    text: String,
}

#[pymethods]
impl PyResponse {
    #[new] // コンストラクタ
    fn new(text: String) -> Self {
        Self { text }
    }
}

#[pymodule]
fn mymodule(module: &Bound<'_, PyModule>) -> PyResult<()> {
    module.add_class::<PyGreeting>()?;
    module.add_class::<PyResponse>()?;
    Ok(())
}

このあたりはPythonスクリプトとRust両方から使用できる型を定義し、モジュールにまとめている箇所です。 #[pyclass(name = "...")] 属性で型を定義しつつ、Pythonから違う名前で見えるようリネームしています。 #[pymethods] でコンストラクタ等のメソッドも定義しており、これらもPythonから呼び出せるようになっています。

    pyo3::append_to_inittab!(mymodule);
    pyo3::prepare_freethreaded_python();

Pythonの初期化を行います。 append_to_inittab!() で自作モジュールが読み込まれるように登録し、 prepare_freethreaded_python() で初期化します。なお、pyo3 に auto-initialize という feature を指定していれば prepare_freethreaded_python() は不要になります。

    let script = r#"
import mymodule

def greeting(g):
    if g.text == "hello":
        res = "こんにちは"
    else:
        res = "なんですか?"
        
    return mymodule.Response(res)
"#;

実行するPythonスクリプトを定義します。greeting 関数は Greeting 型の引数を受け取り、 Response 型を返します。この例では上記で用意したカスタムの型を受け渡ししていますが、文字列のようなプリミティブな型もやり取りできます。

    println!(
        "{:?}",
        run_script(
            script,
            PyGreeting {
                text: "hello".into()
            }
        )
    );
    println!(
        "{:?}",
        run_script(
            script,
            PyGreeting {
                text: "hoge".into()
            }
        )
    );

スクリプトを実行し、その結果を表示します。

fn run_script(script: &str, greeting: PyGreeting) -> PyResponse {
    let result: PyResult<PyResponse> = Python::with_gil(|py| {
        let pyfunc: Py<PyAny> = PyModule::from_code_bound(py, &script, "", "")?
            .getattr("greeting")?
            .into();

スクリプトを実際に実行します。Pythonスクリプトの実行は、Python::with_gil に渡した関数の中で行います。 PyModule::from_code_bound でPythonのソースコードを読み込み、その中の greeting 関数を取得します。

        let result = pyfunc.call1(py, (greeting,));
        let response = match result {
            Ok(response) => response,
            Err(e) => {
                eprintln!("{}", e);
                if let Some(e) = e.traceback_bound(py) {
                    if let Ok(e) = e.into_gil_ref().format() {
                        eprintln!("{}", e); // スタックトレースの表示
                    }
                }
                panic!();
            }
        };
        let response: PyResponse = response.extract(py)?;
        Ok(response)

greeting 関数を実行し、その結果を引き出します。Python関数内でエラーが起きた場合スタックトレースの情報が含まれていれば、 e.traceback_bound(py) で取り出して表示します。最後に extract() で返り値を PyResponse に変換して結果を返します。

実行結果

実行してみると以下のようになります。

PyResponse { text: "こんにちは" }
PyResponse { text: "なんですか?" }

Pythonスクリプト内に記述した分岐の通りの返答を返してくれています。

マルチスレッドの場合

上記の例では関係ないのですが、実際にPyO3を利用して実装した時に苦労した点として、with_gil に渡した関数を実行できるのは1スレッドという制限があげられます。言い換えればPythonスクリプトを実行できるのは常に1つのスレッドのみです。これはPythonのGILという仕組みによるもののようです。

allow_threads を使えば、一時的にGILを解放できるので、Rustのコード中でブロックするような処理があるところはこれを追加して解決しました。Pythonソースコード内にブロックするところがあると処理が進まなくなる可能性は出るので、完全な解決ではありませんが。

まとめ

今回はRust製アプリケーションにPythonを組み込む方法を紹介しました。ビルド不要で柔軟に挙動を変えたいところ、たとえばデータのフィルタリング処理やゲームのスクリプト等があり、メインの部分はRustで書いているアプリケーションがあるなら、今回の手法は役立つと思います。