RustでPythonの拡張モジュールを書く (setuptools-rust)

今日は setuptools-rust を使い、 Rust 言語で Python 拡張モジュールを勉強がてら書いてみたので記録しておく。 1.0 未到達なので、すぐに変わってしまうかもしれないけれど、 大きな流れは変わらないだろう…と期待している。

• 2021-11-28: setuptools-rust 1.0.0で TOML パース用の依存パッケージが toml から tomli に変更されたことを受けて微修正
• 2021-12-09: setuptools-rust 1.1.2 で TOML パース用の依存パッケージ tomli に依存しないよう変更されたことを受けて微修正

動作確認環境

• OS
  • Windows 10
• Python
  • CPython 3.8.9
  • setuptools-rust 0.12.1
• Rust
  • rustc 1.52.1
  • cargo 1.52.0
  • stable-x86_64-pc-windows-msvc
  • pyo3 0.13.2

お勉強のゴール（作るモノ）

Python で次のように使える、ただし Rust で書かれた Python の拡張モジュール mylib.greetings を作る:

>>> import mylib.greetings
>>> mylib.greetings.hello("George")
Hello, George!

2021-05-29 現在、いくつかの方法があるけれど、その中で今日は setuptools-rust を試してみることにした。最終的には次のような成果物が出来上がる:

1. mylib という名前の Rust の crate
2. mylib という名前の Python のパッケージ

なお、これは setuptools のプラグインなのだけれど、 setuptools の深淵を覗き込まずともセットアップできるよう書いたつもり。

手順概要

1. Rust プロジェクトを作成
2. setuptools-rust 用に数個のファイルを作成
3. Rust 側で、Python 側に公開するコードを実装する
4. ビルドとテスト
5. 使う

1. Rust プロジェクトを作成

普通に Cargo コマンドでプロジェクトの雛型を作成する。 これ自体を実行することは無いだろうから、ライブラリとして作れば良いだろう。

PS ~\src\mylib> cargo new --lib mylb
     Created library `mylib` package

2. setuptools-rust 用に数個のファイルを作成

setuptools-rust オフィシャルのREADMEに従って、 以下のファイルを作っておく。

MANIFEST.in

このファイルでは Python のパッケージを作る場合と同様に、 sdist （source distribution, 多くの場合 .tar.gz） に同梱すべきファイルを指定する。 Rust で拡張モジュールを作る場合は基本的に Rust のソースコード類をこれで指定して、sdist に含めるよう設定する。 さもないと、Python 的には Cargo.toml やら .rs ファイルやらは意味の無いファイルと見えてしまうため、 それらを含まない sdist が作られてしまう。 すると、配布先の環境でいざパッケージをビルドしようとしたときに Rust のソースコードが「無い」ため、ビルドエラーになってしまう。 （パッケージのビルドと聞いてピンと来なければ、Python の sdist と binary wheel などといったキーワードで検索されたい）

ということで、オフィシャル README に倣って Cargo.toml と src 下の全ファイル（≒Rust のソースファイル）を指定する例を以下に掲載:

include Cargo.toml
recursive-include src *

pyproject.toml

比較的新しい Python プロジェクトの定義ファイル [PEP 518] では、 そのプロジェクトのビルドに使うツールを指定できる。 これに、setuptools-rust を使う設定を記入する。

[build-system]
requires = ["setuptools", "wheel", "setuptools-rust"]

setup.py

比較的新しい setuptools を使った Python パッケージでは、 pyproject.toml と、setup.py または setup.cfg を組み合わせてプロジェクトの情報を定義する。 setuptools-rust は setuptools のプラグインなので、 このあたりは同じ…と言いたいところだけれど、 「Rust で書いたライブラリをどのように Python 側に見せるのか」 を宣言する都合上、setup.cfg ではなく setup.py を使う必要がある （っぽい）。

具体例を書いてみる（オフィシャル README のものを多少変えている）:

from setuptools import find_packages, setup
from setuptools_rust import Binding, RustExtension

setup(
    name="mylib",
    version="0.1.0",
    rust_extensions=[
        RustExtension(
            "mylib.greetings",  # Make it importable with this name
        )
    ],
    packages=find_packages(),
    # rust extensions are not zip safe, just like C-extensions.
    zip_safe=False,
)

setup 関数の rust_extensions キーワード引数に RustExtension オブジェクトを渡している。 オブジェクト生成時、第 1 引数には Python コード上で、何という名前で Rust で書いたモジュールを import できるようにするかを指定する。 上記の例では mylib.greetings という名前を指定している （後で Rust 製のライブラリを作成するときに重要なので覚えていよう）。 他にも指定できるオプションがあるので、興味あれば setuptools-rust の API リファレンスを参照されたい。

setup 関数には rust_extensions 以外のキーワード引数も指定している。 これらは Python パッケージとしての名前・バージョン番号などとして使われる。 今回は Rust の crate と Python のパッケージを同じ名前・バージョン番号にする方針なので、setup.py の中で Cargo.toml からパッケージ名とバージョン番号を読み出して流用している。なお共通化すべきかどうかは個別ニーズ次第なので、これがベストプラクティスというわけではないので注意されたい。 ちなみに、setuptools-rust （v1.0.0 以降）の依存パッケージには tomli が含まれているので tomli パッケージは setup.py の中で確実に使える。

3. Python に公開するコードを Rust で実装する

Python パッケージ側の準備は整ったので、今度は Rust 側でライブラリの実装を行う。

まず最初に Cargo.toml に PyO3 を使って Python コードと相互運用するための設定を追加する。具体的には、lib セクションの crate-type で cdylib を追加する（C 向け動的リンクライブラリを作る指示）。また dependencies セクションに pyo3 を加え、features フラグに extension-module を加える。必要部分を抜粋すると以下のようになる:

[lib]
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "~0.13.2", features = ["extension-module"] }

続いて、ライブラリのロジックを書く。たとえば次のような感じになる:

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

#[pyfunction]
/// Print a greeting message and return length of it.
fn hello(whom: Option<&str>) -> usize {
    let msg = if let Some(whom) = whom {
        format!("Hello, {}!", whom)
    } else {
        format!("Hello, extension in Rust!")
    };
    println!("{}", msg);
    msg.len()
}

#[pymodule]
/// The module containing a greeting function "hello".
fn greetings(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    // This function defines what the module has.
    // Name of this must match the "target" specified to RustExtension in setup.py.

    // Expose a function to Python side
    m.add_wrapped(wrap_pyfunction!(hello))?;

    Ok(())
}

// Unit tests
#[cfg(test)]
mod test {
    use super::hello;

    #[test]
    fn hello_some() {
        assert_eq!(hello(Some("John")), 12);
    }

    #[test]
    fn hello_none() {
        assert_eq!(hello(None), 25);
    }
}

ポイントを簡単に説明すると:

• Python から関数として呼び出したい hello 関数に、 pyo3 の pyfunction マクロを付けている
• Python から import されるモジュールに何が含まれているのかを 定義する関数 greetings を、pymodule マクロを付けている
  • setup.py で「mylib.greetings という名前でインポートする」 と指定したのに合わせて greetings という名前を使用
  • 名前が違うと import 時にエラーが起こる
• greetings 関数で hello 関数をモジュールに登録するとき、 "wrap" している
  • 詳細未調査だけれど、おそらく引数・戻り値を Rust の型と Python の型との間で相互変換していると思う （str と PyUnicode の相互変換、など）
• #[cfg(test)] 以下は単体テスト
  • 本質的でないので Rust 未経験者は無視してください

念のため hello 関数の内容を簡単に書いておくと、引数 name に値が指定されたなら「Hello, その名前」を、さもなくば (i.e.: name = None) 「Hello, extension in Rust!」を表示して、その文字数を返す内容となっている。

4. ビルドとテスト

Python 側も Rust 側も準備が整ったので、ビルド・テストを試してみる。

Rust 側

なお、概念的に Python の mylib パッケージは Rust の mylib crate に依存している。この構造なので、後者は前者とは独立してビルド・テストができる。試しにテスト（同時にビルドも）を走らせてみる:

PS ~\src\mylib> cargo test
～略～
running 2 tests
test test::hello_none ... ok
test test::hello_some ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s

確かに単独でビルドでき、自動テストも実行できる。

Python 側

Python 側は、少し込み入っている。好きな方法を使うと良いかと:

• 配布用パッケージ (sdist, bdist_wheel) を作る
  1. PyPA build でビルドする
     • pip install build した環境で、
     • python -m build
  2. setuptools でビルドする
     • pip install setuptools-rust した環境で、
     • python setup.py sdist bdist_wheel
• 編集可能モードで、現在の環境にインストールする
  1. pip で editable mode で install する
     • pip install -e /path/to/project
  2. setuptools で develop mode インストールを行う
     • pip install setuptools-rust した環境で、
     • python setup.py develop

配布用パッケージを作成する場合、dist ディレクトリ下に .tar.gz ファイル (sdist) と .whl ファイル (bdist_wheel) が作成されていれば OK （心配なら pip install そのファイル してみる）。 編集可能モードでのインストールの方は、普通に pip show の結果、 開発プロジェクトがあるディレクトリパスが Location 値に表示されれば OK:

PS ~\src\mylib> pip show mylib
Name: mylib
Version: 0.1.0
Summary: UNKNOWN
Home-page: UNKNOWN
Author: None
Author-email: None
License: UNKNOWN
Location: c:\users\sgryjp\src\mylib  # site-package ではない
Requires:
Required-by:

5. 使う

配布用パッケージを pip install するか、編集可能モードでインストールした環境で Python インタプリタを起動して使ってみる:

>>> import mylib.greetings
>>> n = mylib.greetings.hello()
Hello, extension in Rust!
>>> n
25
>>> n = mylib.greetings.hello("George")
Hello, George!
>>> n
14

なお、Rust の関数に付けたドキュメントコメントは Python の docstring として使われる:

>>> import mylib.greetings
>>>
>>> help(mylib.greetings.hello)
Help on built-in function hello:

hello(...)
    Print greeting message.

>>>
>>> help(mylib.greetings)
Help on module mylib.greetings in mylib:

NAME
    mylib.greetings - The module containing some greeting functions.

FUNCTIONS
    hello(...)
        Print greeting message.

DATA
    __all__ = ['__doc__', 'hello']

FILE
    c:\users\sgryjp\src\mylib\mylib\greetings.cp38-win_amd64.pyd

以上。