📝

RustでPyO3+maturin環境の非破壊stubgenerator

に公開

はじめに

久しぶりの投稿になります。
今回は、RustのPyO3で使われるアトリビュートを解析し、既存のコードを壊さずにPythonのスタブファイル(.pyi)を生成するサブコマンドを作成したお話を紹介します。

作成の経緯

PythonでGUIアプリケーションを開発する中で、MVVMアーキテクチャを採用し、ViewとViewModelはPythonで、Model部分をRustで実装しようと考えました。
PyO3とMaturinを使ってRustコードをPythonにバインドする際、Rust関数の変更がPython側からもすぐに把握できるようにしたくなりました。

そのため、Rustの関数や型の定義をもとにスタブファイルを自動生成し、Pythonからも補完や型エラーがわかりやすくなるような仕組みを整える必要がありました。
このニーズから、非破壊的にアトリビュートを解析してスタブファイルを生成するツールの開発を始めました。

既存ツールとの比較と採用理由

調査の中で、以下のような既存ツールも見つかりました:

pyo3-stub-gen
PyO3のコードからスタブファイルを生成できるツールです。
しかし、このツールでは関数や構造体に対して専用のアトリビュート(例:#[gen_stub_pyfunction] など)を明示的に追加する必要があり、アトリビュートが煩雑になる点が気になったため、今回は見送りました。

setuptools-rust
Pythonのビルドシステム(setuptools)とRustを連携させるための公式ツールです。
ただし、これはPythonパッケージとしてのセットアップ時に組み込む必要があり、柔軟性に欠けると感じました。

今回の目的は、既存のRustプロジェクトに後からでも導入でき、かつ非破壊的にスタブファイルを生成できる仕組みを作ることでした。
そのため、Cargoのサブコマンドとして動作する形で実装することにしました。
これにより、PyO3のコードを書いている途中でも、手軽にスタブファイルを生成できるようになりました。

機能概要

現在対応している内容:

  • 関数の定義(引数・戻り値を含む)
  • OptionやVecの型の展開
  • Rustの基本型の変換(int, float, str など)
  • uv ワークスペースのサポート

今後実装予定:

  • Rust構造体からPythonのclass定義へ
  • numpy配列やTuple型の変換
  • Pythonの型コメント形式への完成度推進

使い方

基本的な使い方

cargo pystubgen

スタブ変換例

このラストのコードを

test_code.rs
/// ハッシュマップを使用した関数
#[pyfunction]
pub fn test_hashmap_types(
    map: HashMap<String, i32>,
) -> PyResult<HashMap<String, i32>> {
    let mut result = HashMap::new();
    for (key, value) in map {
        result.insert(format!("new_{}", key), value * 2);
    }
    Ok(result)
}

このようなスタブに変換する

test.pyi
#  ハッシュマップを使用した関数
def test_hashmap_types(map: dict[str, int]) -> dict[str, int]:

実装の工夫

  • Rustの構文解析に [syn] クレートを使用し
perser.rs
pub fn parse_function_data(item: &syn::ItemFn) -> RustFunctionData{
    RustFunctionData{
        name: parse_function_name(item),
        args: parse_function_args(item),
        return_type: parse_function_return_type(item),
        attributes: parse_function_attributes(item),
        doc: parse_function_doc(item),
    }
}

pub fn parse_function_name(item: &syn::ItemFn) -> String{
    item.sig.ident.to_string()
}

pub fn parse_function_args(item: &syn::ItemFn) -> Vec<(String, String)>{
    item.sig.inputs.iter().map(|arg|match arg{
        syn::FnArg::Typed(arg) => (arg.pat.to_token_stream().to_string(), arg.ty.to_token_stream().to_string()),
        _ => ("self".to_string(), "self".to_string()),
    })
    .collect()
}

pub fn parse_function_return_type(item: &syn::ItemFn) -> String{
    match &item.sig.output {
        syn::ReturnType::Default => "None".to_string(),
        syn::ReturnType::Type(_, ty) => ty.to_token_stream().to_string(),
    }
}

pub fn parse_function_attributes(item: &syn::ItemFn) -> Vec<String> {
    item.attrs.iter()
        .filter_map(|attr| match &attr.meta {
            Meta::Path(path) if !attr.path().is_ident("doc") => {
                Some(path.segments.last().unwrap().ident.to_string())
            }
            _ => None,
        })
        .collect()
}


pub fn parse_function_doc(item: &syn::ItemFn) -> String{
    let mut doc = String::new();
    item.attrs.iter().for_each(|attr|{
        match &attr.meta {
            Meta::NameValue(name_value) =>match &name_value.value {
                syn::Expr::Lit(lit_str) => {
                    let doc_str = match &lit_str.lit {
                        syn::Lit::Str(lit_str) => lit_str.value(),
                        _ => "".to_string(),
                    };
                    doc.push_str(&doc_str);
                }
                _ => {}
            }
            Meta::Path(_) => {},
            Meta::List(_) => {},
        }
    });
    doc
}

このようにパースし、Anarayzaでpythonの型に変換します。

今後の展望

  • Rust構造体 → Pythonクラスの自動変換
  • maturin や build.rs に組み込んでビルド時に自動出力

まとめ

RustとPythonのハイブリッド開発は、速度と柔軟性の両立にとても有効です。cargo pystubgen は、PyO3を使った開発において、Python側での補完や型エラー検出を支援する生産性向上ツールです。

興味があれば、GitHubのStarやIssueなどもお待ちしています!

リンク

https://github.com/GamingChikuwabu/cargo-pystubgen.git
https://crates.io/crates/cargo-pystubgen

Discussion