🎇

uv + maturin + PyO3でRustをPythonから呼ぶ

に公開

PyO3/maturin は、Rust コードを Python から使えるようにバインディングを生成し、Pythonパッケージとしてビルド・配布するためのツールチェインです。
この記事では、maturin を用いて Rust ライブラリを Python から使用し、更に Rust 実装に基づき Python の型定義ファイル(.pyi)を自動生成したり、独自の Python 実装部分を記述する方法について説明します。

準備

Python のモダンなパッケージマネージャである uv をインストールします。既にインストール済みの方は飛ばしてください。
他のパッケージマネージャを使いたいという方は適宜読み替えてください。

https://docs.astral.sh/uv/getting-started/installation/

maturin の基本的な使い方

maturin new でプロジェクトを生成します。

$ maturin new -b pyo3 maturin_sample
  ✨ Done! New project created maturin_sample
$ tree maturin_sample
maturin_sample
├── Cargo.toml
├── pyproject.toml
└── src
    └── lib.rs

仮想環境(venv)をアクティベートします。このコマンドだけだと仮想環境に入れない場合もあるので、出力に従い、source .venv/bin/activate などを実行してください。

uv venv

maturin develop で開発用、 maturin build でリリース用のビルドを行います。いったん適当にビルドしてみましょう。

uvx maturin develop

これでもう Python から maturin_sample が呼べます。

tmp.py
import maturin_sample
print(maturin_sample.sum_as_string(1, 2))

として、

$ uv run tmp.py
3

いい感じですね。

なお、この際ビルドされた wheel は target/wheels/maturin_sample-~~~.whl に配置されます。
ローカルの別プロジェクトから手軽に試すだけなら、uv add /path/to/maturin_sample/target/wheels/maturin_sample-~~~.whl などで直接依存を追加でき、便利です。

型定義を自動生成する

stub ファイルとは?

Python から maturin_sample が呼べましたが、1 つ大きな問題点があります。
この maturin_sample は単なる .so ファイルとして、インタプリタにより完全に動的にロードされています。すなわち、型やモジュール定義に関する静的な情報が失われてしまい、型検査やIDEによる補完が不可能な状態になっています。

試しに、最新の Python の型検査器である ty で確認してみましょう。ty は記事執筆時点ではまだプレリリースですが、適当に使うぶんにはいい感じです。(ty が嫌/ダメそうな場合は、mypy や pyright でも良いです)
tmp.py を型エラーを含む状態にします。

tmp.py
import maturin_sample
print(maturin_sample.sum_as_string(1, 2))
+ print(maturin_sample.sum_as_string(1, "2"))
$ uvx ty check tmp.py
error[unresolved-attribute]: Type `<module 'maturin_sample'>` has no attribute `sum_as_string`
 --> tmp.py:3:7
  |
1 | import maturin_sample
2 | print(maturin_sample.sum_as_string(1, 2))
3 | print(maturin_sample.sum_as_string(1, "2"))
  |       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  |
info: `unresolved-attribute` is enabled by default

そもそも maturin_sample に sum_as_string という関数がないよ、と言っていますね。
前述したとおり、モジュール定義に関する静的な情報が失われているので、型検査器は関数の存在を知れない、というわけです。

そこで、PEP561 で定義される stub file (.pyi)と呼ばれる型定義ファイルを別に用意します。js/ts における .d.ts ファイルのようなものです。

pyo3-stub-gen で自動生成する

stub file を頑張って手動で書いても良いのですが、pyo3-stub-gen というツールで自動生成が可能です。これは試験的/不完全なツールですが、労力は大幅に削減できますし、簡単なプロジェクトには十分です。

最初に、現在の lib.rs の実装を確認しておきましょう。

lib.rs
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 maturin_sample(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

まず、依存を追加します。

cargo add pyo3-stub-gen

次に、現在の Rust 側実装における pyfunctionpyclass の定義に、それぞれ gen_stub_pyfunctiongen_stub_pyclass を付けていきます。
あと、lib.rsdefine_stub_info_gatherer! を追記します。

lib.rs
use pyo3::prelude::*;
+ use pyo3_stub_gen::{define_stub_info_gatherer, derive::gen_stub_pyfunction};

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

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

+ define_stub_info_gatherer!(stub_info);

最後に、Cargo.toml を編集して crate-typerlib を追加します。

Cargo.toml
[lib]
name = "maturin_sample"
+ crate-type = ["cdylib", "rlib"]
- crate-type = ["cdylib"]

これで Rust 実装側の準備は完了です。


続いて、型定義を自動生成するプロセス自体のバイナリを Rust クレート内に新たに生やしてあげます。

mkdir ./src/bin
touch ./src/bin/stub_gen.rs
stub_gen.rs
use pyo3_stub_gen::Result;

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

最後に、この stub_gen を実行してあげます。

cargo run --bin stub_gen

これで maturin_sample.pyi という型定義ファイルが生成されます。

$ cat maturin_sample.pyi
# This file is automatically generated by pyo3_stub_gen
# ruff: noqa: E501, F401

import builtins

def sum_as_string(a:builtins.int, b:builtins.int) -> builtins.str:    r"""
    Formats the sum of two numbers as string.
    """

かなり汚いコードなので、気になる場合は手動で綺麗にしてあげましょう。

maturin_sample.pyi が生成されたこの状態で maturin develop などを実行してあげると、maturin が自動でこの型定義ファイルを収集し、パッケージに含めてくれます。ただ、個人的には以下の方法で明示的に管理してあげるのがおすすめです。


いずれにせよ、これで型定義が生成できたので、型検査器が正しく動くか確認してみましょう。

$ maturin develop
$ uvx ty check tmp.py
error[invalid-argument-type]: Argument to function `sum_as_string` is incorrect
 --> tmp.py:3:39
  |
1 | import maturin_sample
2 | print(maturin_sample.sum_as_string(1, 2))
3 | print(maturin_sample.sum_as_string(1, "2"))
  |                                       ^^^ Expected `int`, found `Literal["2"]`
  |
info: Function defined here
 --> maturin_sample.pyi:6:5
  |
4 | import builtins
5 |
6 | def sum_as_string(a:builtins.int, b:builtins.int) -> builtins.str:    r"""
  |     ^^^^^^^^^^^^^                 -------------- Parameter declared here
7 |     Formats the sum of two numbers as string.
8 |     """
  |
info: `invalid-argument-type` is enabled by default

しっかりと「引数の型が間違ってるよ」と指摘してくれています。ご丁寧に関数の定義位置まで教えてくれました。いい感じですね!

独自の Python 実装部分を記述する

Rust 実装をそのまま Python から実行することは可能になりましたが、Python パッケージとして配布する際には、Python 特有の実装部分を追加したいことがあると思います。また、先程のように型定義ファイルを扱いたい場合なども、明示的に Python パッケージ部分を管理できると便利です。

まず、Python パッケージ部分を格納する python ディレクトリを生やし、その存在を Cargo.toml に記述します。このディレクトリ名は任意です。

mkdir -p python/maturin_sample
Cargo.toml
+ [tool.maturin]
+ python-source = "python"

中に Python 独自の実装部分や型定義ファイルを置きます。この書き方は通常の Python パッケージと同じです。
Rust 実装のライブラリ(.so)部分は python/maturin_sample/maturin_sample.~~~.so あたりに置かれることになり、import .maturin_sample などで import できます。
それぞれこんな感じにします。


__init__.py: import するためのエントリポイント。Rust 実装のライブラリ(.so) に丸投げするだけ。Ruff(linter) によるエラーを抑制する指示を入れています

__init__.py
# ruff: noqa: F403, F405

from .maturin_sample import *

__doc__ = maturin_sample.__doc__
if hasattr(maturin_sample, "__all__"):
    __all__ = maturin_sample.__all__

__init__.pyi: 型定義。前述の自動生成した maturin_sample.pyi をここに入れます

__init__.pyi
import builtins

def sum_as_string(a:builtins.int, b:builtins.int) -> builtins.str:    r"""
    Formats the sum of two numbers as string.
    """

py.typed: 型付きライブラリであると型チェッカに認識させるための空のフラグファイル。内容は何もないが、必須

py.typed

これでビルドすれば、python/maturin_sample の中身に Rust 実装部分が加えられた上で、 Python パッケージとしてビルドされます。
Python 独自の関数や型定義を追加したい場合は、通常の Python パッケージと同じ感覚でここに追加していくだけでOKです。

型定義ファイルの配置方法として、maturin_sample.pyi を置く方法と、__init__.pyi を置く方法の2通りあり、これらは二者択一であることに注意してください。個人的には後者がおすすめです。

おわりに

本記事は、以下の記事を大変参考にさせていただきました。素晴らしい情報を公開してくださった著者の方々に、この場を借りて心より感謝申し上げます。
https://zenn.dev/jij_inc/articles/pyo3-mannually-type-stub-file

また、パッケージの開発とローカルでの簡単なお試しについては本記事で説明しましたが、実際の配布の手順については公式ドキュメントなどを参考にしてください。

Discussion