Rust pyo3-stub-genメモ

Rustのpyo3のpyo3-stub-genクレートのメモ。

概要

pyo3-stub-genは、Pythonの型スタブファイル（拡張子pyiのファイル）を生成するクレート。

型スタブファイルとは、Pythonの関数やクラス・メソッド・プロパティーの型定義（型ヒント）のみを記述したファイル。
IDEの補完機能などで使われる。

なお、pyo3-stub-genはPython 3.10以降が対象らしい。（それ以前向けのPyO3には対応していないらしい）

pyo3-stub-genでは、PyO3を使って定義された関数やメソッドから、型スタブファイルを生成できる。

生成対象のpyclass・pymethods・pyfunctionに、それぞれgen_stub_pyclass・gen_stub_pymethods・gen_stub_pyfunction属性を付ける必要があるので、ちょっと面倒だが。
（それらの属性が付いていない関数やメソッドは、生成対象にならない）

use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods, gen_stub_pyfunction};

基本的にはRustの型からPythonの型が自動的に判別されるが、判別できない場合は手動で指定することも出来る。

また、「///」によって書かれたRustのドキュメンテーションコメントも型ヒントに転写される。

そして、生成用のスクリプトを書いて実行すると、型スタブファイルが生成される。

例

pyo3-stub-genを使って型スタブファイルを生成する例。

先に生成スクリプトを作っておきたいところだが、本体コードに手を入れてからでないとコンパイルエラーになるので、後回しにする。

Cargo.tomlの定義

pyo3-stub-gen自体は普通にdependenciesに追加すればいいが、crate-typeも修正する必要がある。

example-pyo3/Cargo.toml：

〜
[lib]
name = "example_pyo3"
crate-type = ["cdylib", "rlib"]

〜
[dependencies]
pyo3 = { version = "0.27.2" }
pyo3-stub-gen = "0.18.0"

通常のPyO3プロジェクトはcrate-typeにcdylibのみ指定されているが、型スタブファイルの生成スクリプトをコンパイル・実行するために、rlibも加えておく必要がある。

生成対象の指定

まず、型スタブを生成する対象となる関数やメソッドに、生成対象であることを示す属性を付ける。

use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods, gen_stub_pyfunction};

example-pyo3/src/lib.rs：

use pyo3::prelude::*;

mod my_class;

#[pymodule]
mod example_pyo3 {
    use pyo3::{prelude::*, types::*};
    use pyo3_stub_gen::derive::gen_stub_pyfunction;

    #[gen_stub_pyfunction]
    #[pyfunction]
    fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
        let r = format!("pyo3==>{}", a + b);
        Ok(r)
    }

    #[pymodule_export]
    use super::my_class::MyClass;
}

example-pyo3/src/my_class.rs：

use pyo3::prelude::*;
use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};

#[gen_stub_pyclass]
#[pyclass]
pub struct MyClass {
    #[pyo3(get, set)]
    value: String,
}

#[gen_stub_pymethods]
#[pymethods]
impl MyClass {
    #[new]
    fn new(value: String) -> Self {
        Self {
            value,
        }
    }

    fn my_method(&self) -> &str {
        &self.value
    }
}

スタブ情報収集メソッドの定義

pyo3-stub-gen（の生成スクリプト）が型情報を収集するためのメソッドを用意する必要がある。

そのメソッドはdefine_stub_info_gatherer!マクロで定義できる。

example-pyo3/src/lib.rs：

use pyo3::prelude::*;
use pyo3_stub_gen::define_stub_info_gatherer;

#[pymodule]
mod example_pyo3 {
    〜
}

define_stub_info_gatherer!(stub_info);

引数の「stub_info」は、マクロによって生成されるメソッドのメソッド名。
生成スクリプトからこのメソッドを呼び出す。

（メソッド名は何でもいいが、慣例的にstub_infoにしておくのが良いと思う）

生成スクリプト

型スタブファイルを生成するスクリプトを用意する。

example-pyo3/src/bin/stub_gen.rs：

use pyo3_stub_gen::Result;

fn main() -> Result<()> {
    let stub = example_pyo3::stub_info()?;
    stub.generate()?;
    Ok(())
}

自分のモジュール（この例ではexample_pyo3）の本体コード上に定義したスタブ情報収集メソッド（この例ではstub_info()）を呼び出している。

※Cargo.tomlのcrate_typeにrlibを追加していないと、自分のモジュール名を指定した部分がコンパイルエラーになるので注意。

このスクリプトを実行すると、型スタブファイル（拡張子pyiのファイル）が生成される。

$ cd example-pyo3
$ cargo run --bin stub_gen

$ ls *.pyi
example_pyo3.pyi

なお、生成対象が無い（gen_stub_pyXXX属性がひとつも付けられていない）と、型スタブファイルは生成されない。

ちなみに、「uv run maturin build」コマンドによってwheelファイル（配布用ファイル。拡張子whl）を生成すると、
置かれている型スタブファイルがwheelファイル内に「__init__.pyi」というファイル名で自動的に入る。

Decimalの例

Pythonのdecimal.Decimalは、PyO3ではrust_decimal::Decimalで扱うことが出来る。

    #[gen_stub_pyfunction]
    #[pyfunction]
    fn inc_decimal(value: rust_decimal::Decimal) -> PyResult<rust_decimal::Decimal> {
        let value = value + rust_decimal::Decimal::ONE;
        Ok(value)
    }

しかしpyo3-stub-genは、デフォルトではrust_decimalに対応してないので、コンパイルエラーになる。

error[E0277]: the trait bound `rust_decimal::Decimal: PyStubType` is not satisfied

pyo3-stub-genのfeaturesにrust_decimalを指定すると、rust_decimal::Decimalが認識されるようになる。

example-pyo3/Cargo.toml：

〜
[dependencies]
pyo3 = { version = "0.27.2", features = ["rust_decimal"] }
pyo3-stub-gen = { version = "0.18.0", features = ["rust_decimal"] }
rust_decimal = "1.40.0"

手動で型ヒントを指定する例

型がpyo3-stub-genで自動判別できない場合（コンパイルエラーになる場合）は、手動で指定する。
あるいは、PyAnyを指定している型に対して、詳細な型を手動で指定できる。

例としてDecimalを指定してみる。
（rust_decimalはpyo3-stub-genが対応しているが、pyo3-stub-genにrust_decimalフィーチャーを指定せず、手動で型ヒントを指定することも出来る）

gen_stub_pyfunctionの場合

    #[gen_stub_pyfunction(python = r#"
        def inc_decimal(value: decimal.Decimal) -> decimal.Decimal: ...
    "#)]
    #[pyfunction]
    fn inc_decimal(value: rust_decimal::Decimal) -> PyResult<rust_decimal::Decimal> {
        let value = value + rust_decimal::Decimal::ONE;
        Ok(value)
    }

gen_stub_pyfunction属性のpythonで（関数全体の）型ヒントを記述する。
（ここに書いた内容が型スタブファイルにそのまま出力される）

「r#"〜"#」はRustの生文字リテラル構文で、「"」や「\」の文字を含むことが出来る。（上記の例では使ってないが＾＾；）

gen_stub_pymethodsの場合

gen_stub_pymethods内のメソッドにはgen_stub_pyfunction属性を付けられないので、代わりにgen_stub属性を使用する。
（gen_stub_pyfunction属性が付いている関数でもgen_stub属性は使用できる模様）

use pyo3::prelude::*;
use pyo3_stub_gen::derive::*;

#[gen_stub_pyclass]
#[pyclass]
pub struct MyDecimal {
    value: Option<rust_decimal::Decimal>,
}

#[gen_stub_pymethods]
#[pymethods]
impl MyDecimal {
    #[new]
    fn new(
        #[gen_stub(override_type(type_repr = "decimal.Decimal | None", imports=("decimal")))]
        value: Option<rust_decimal::Decimal>,
    ) -> PyResult<Self> {
        Ok(Self { value })
    }

    #[getter]
    #[gen_stub(override_return_type(type_repr = "decimal.Decimal | None", imports=("decimal")))]
    fn get_value(&self) -> &Option<rust_decimal::Decimal> {
        &self.value
    }
}

メソッドの引数の型ヒントを指定する場合は、引数にgen_stub属性を付けて、override_typeでPythonの型（型ヒント）を指定する。

メソッドの戻り値の型ヒントを指定する場合は、メソッドにgen_stub属性を付け て、override_return_typeでPythonの型（型ヒント）を指定する。

importsを書くと、型スタブファイルの冒頭にimport文が生成される。

※gen_stubを使う場合、個別にインポート（use）しなくてもいい模様。（gen_stub_pyclassと）gen_stub_pymethodsをインポートしていれば使えるっぽい。どうしてそれでいいのか、謎。 （gen_stub_pymethodsの中でgen_stubマクロを定義しているのか？）

独自クラスの例

gen_stub_pyclass属性が付けられた独自クラス（やgen_stub_pyclass_enum属性が付けられた列挙型）を引数や戻り値で使っている場合は、pyo3-stub-genが自動的に判別してくれる。

そうでないクラス（pyclass）はデフォルトではpyo3-stub-genで判別できないので、gen_stub属性を使う。

    #[gen_stub_pyfunction]
    #[pyfunction]
    #[gen_stub(override_return_type(type_repr = "example_pyo3.MyClass", imports=("example_pyo3")))]
    fn create_my_class(value: String) -> MyClass {
        MyClass { value }
    }

オーバーロードの例

引数や戻り値の型に複数のパターンがあるときは、型ヒントをオーバーロードする。[2026-02-10]

    use pyo3_stub_gen::derive::*;
    use pyo3_stub_gen::inventory::submit;

    #[gen_stub_pyfunction] // ひとつめの型ヒント
    #[pyfunction]
    fn overload_example(arg: Bound<PyAny>) {
        println!("overload_example(): arg={}", arg);
    }

    // ふたつめの型ヒント
    submit! {
        gen_function_from_python! {
            r#"
            @overload
            def overload_example(arg: str) -> None: ...
            "#
        }
    }