Zenn
👨‍🦱

[PyO3入門]Rustで自作したTOMLパーサーをPythonから呼び出す

2023/05/24に公開
Python
Rust
toml
#
maturin
#
pyo3
tech

概要

最近、RustでTOMLパーサーを実装してみるという記事を書きました。
この記事では、自作したTOMLパーサーをPythonから呼び出せるようにしてみたいと思います。
(パーサーに限らず、重い処理やバグを起こしやすそうな処理をRustで実装してPythonから呼び出すというケース[1]を最近よく目にするので個人的に気になっていました。)

RustとPythonの連携にはPyO3というRustのライブラリを使います。
このライブラリを使うと、RustのコードをビルドしてPythonのモジュールを作成することが出来ます。
本記事では、PyO3の基本的な使い方からはじめ、最終的にはRust製の自作TOMLパーサーをPythonから呼び出すところまで実装します。

本記事の構成は以下の通りです。

  1. PyO3の基本的な使い方
  2. RustとPythonの型のマッピング
  3. RustとPythonの型のマッピングを自分で定義する
  4. エラーハンドリング
  5. Rustで開発したTOMLパーサーをPythonから呼び出す

前提

本記事を書くにあたって、使用したソフトウェアのバージョンは以下の通りです。

  • Rust 1.66.1
  • PyO3 0.18.3
  • CPython 3.10.8
  • maturin 0.15.1

maturinはRustのコードをビルドしてPythonのモジュールを作成するのに使うツールです。
pip install maturin などでインストールできます。

PyO3の基本的な使い方

以下の手順でPyO3を使ってRustでPythonのモジュールを作ることが出来ます。

  1. PyO3のプロジェクトを作成する
  2. Rustでコードを書く
  3. RustのコードをビルドしてPythonのモジュールを作る

順番に見ていきましょう。

1. PyO3のプロジェクトを作成する

PyO3のプロジェクトを新しく作成する場合は、maturin newコマンドを使うと簡単です。
例えば、pyo3_testという名前でプロジェクトを新しく作成する場合は、以下のようにします。

maturin new pyo3_test --bindings pyo3

すると、cargo newを実行した時のようにpyo3_testというディレクトリが作成されその中に以下の3つのファイルが作成されます。[2]

Cargo.toml
[package]
name = "pyo3_test"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "pyo3_test"
crate-type = ["cdylib"]

[dependencies]
pyo3 = "0.18.3"
pyproject.toml
[build-system]
requires = ["maturin>=0.15,<0.16"]
build-backend = "maturin"

[project]
name = "pyo3_test"
requires-python = ">=3.7"
classifiers = [
    "Programming Language :: Rust",
    "Programming Language :: Python :: Implementation :: CPython",
    "Programming Language :: Python :: Implementation :: PyPy",
]


[tool.maturin]
features = ["pyo3/extension-module"]
src/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 pyo3_test(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

以上でプロジェクトの作成は完了です。

上の例は新しくプロジェクトを作成する場合でしたが、既存のRustプロジェクトをPyO3のプロジェクトにする場合は、同じようにCargo.tomlを編集して、pyproject.tomlを作成すればOKです。

2. Rustでコードを書く

PyO3のドキュメントなどを参考に、Rustでコードを書いていきます。
まずはmaturin newコマンドで作成されたプロジェクトのsrc/lib.rsの内容を見てみましょう。

src/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 pyo3_test(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

pymoduleというマクロを使ってPythonのモジュールを作成し、そこに関数やクラスなどを作って追加していきます。
上の例では、pyfunctionというマクロを使いsum_as_stringという関数を作って、add_functionでpyo3_testというモジュールに追加しています。

このまま上のコードをPythonから呼び出してみてもいいのですが、それだけだと面白くないので試しに自分でも何か関数を作ってみましょう。

src/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())
}

+#[pyfunction]
+fn sum_int_vector(v: Vec<i32>) -> PyResult<i32> {
+    v.iter().sum()
+}

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

sum_int_vectorという関数を追加してみました。
見て分かる通り、Vec<i32>を引数にとり、その要素を全て足し合わせた値を返す関数です。

3. maturin developでPythonのモジュールを作る

コードが書けてしまえばあとは非常に簡単で、以下のコマンドを実行するだけでRustのコードをビルドしてPythonのモジュールを作成した上に自分の環境にインストールまでしてくれます。

maturin develop

モジュールの作成とインストールが完了したので早速試してみましょう。
pyo3_testを呼び出すコードを作成して実行してみます。

run.py
import pyo3_test

s = pyo3_test.sum_as_string(1, 2)
print(f'{s}: {type(s)}')

i32 = pyo3_test.sum_int_vector([1, 2, 3, 4, 5])
print(f'{i32}: {type(i32)}')

実行結果は以下の通りです。期待した通りの結果が得られました。

3: <class 'str'>
15: <class 'int'>

RustとPythonの型のマッピング

上の例を見て気付かれたかもしれませんが、
Rust・Python間でデータを受け渡しする際に以下のような型のマッピングが行われていることが分かります。

  • Pythonから渡したintがRustではusizeとして受け取られている。
  • Pythonから渡したintのlistがRustではVec<i32>として受け取られている。
  • Rustから返したStringがPythonではstrとして受け取られている。
  • Rustから返したi32がPythonではintとして受け取られている。

この型同士のマッピングについては、PyO3のユーザーガイドに表としてまとめられています。
https://pyo3.rs/v0.18.3/conversions/tables

例えば、Pythonからintを渡した場合、Rust側では「Any integer type (i32, u32, usize, etc)」になると書かれています。
従って、Rust側で引数がi32とかu32とかusizeとかになっている関数には、Python側からintの値を渡すことが出来るということになります。

この対応表を見ながら、先ほどよりもう少し複雑な型の値を受け渡してみましょう。

lib.rs
use std::collections::HashMap;
use pyo3::prelude::*;

#[pyfunction]
fn tuple_list_to_dict(v: Vec<(String, i32)>) -> PyResult<HashMap<String, i32>> {
    Ok(v.into_iter().collect())
}

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

tuple_list_to_dictは、引数としてVec<(String, i32)>を受け取り、それをHashMap<String, i32>に変換して返す関数です。
対応表より、Vec<(String, i32)>とHashMap<String, i32>は、Python側ではそれぞれlist[tuple[str, int]]とdict[str, int]に対応すると分かります。

実際に、先ほどと同様にビルドし以下のコードを実行してみると、
list[tuple[str, int]]を受け取ってdict[str, int]を返す関数として期待通りに動作します。

run.py
import pyo3_test

tuple_list = [('a', 1), ('b', 2), ('c', 3)]
dic = pyo3_test.tuple_list_to_dict(tuple_list)
print(f'{dic}: {type(dic)}') # {'b': 2, 'a': 1, 'c': 3}: <class 'dict'>

RustとPythonの型のマッピングを自分で定義する

前節でRustとPythonの型のマッピング表を見ましたが、表に書かれていない型の値を受け渡したい場合もあるでしょう。
その場合、Rust・Python間での型変換のためのトレイトを実装する必要があります。型変換に関するトレイトは以下の2つです。

  • FromPyObject: Python から Rust への型変換 (引数)
  • IntoPy: Rust から Python への型変換 (返り値)

それぞれについて簡単な例を見ていきましょう。

FromPyObject: Python から Rust への型変換

まずは、FromPyObjectトレイトについてです。
FromPyObjectトレイトを普通に実装することも出来そうですが、ユーザーガイドを見た感じderiveマクロを使う方法ばかりが紹介されているので、こちらでもその方法でやってみます。

例えば、Python側で作成したPersonというクラスの変数をRustに渡したいとします。
このPersonクラスは、nameとageという属性を持っていて、それぞれstrとintであるとします。

以下のコードは、Python側からPersonのリストを受け取り、それをHashMap<String, i32>に変換して返すという処理を実装したものです。

lib.rs
use std::collections::HashMap;
use pyo3::prelude::*;

#[derive(FromPyObject)]
struct Person {
    #[pyo3(attribute("name"))]
    name: String,
    #[pyo3(attribute("age"))]
    age: i32,
}

#[pyfunction]
fn person_list_to_dict(v: Vec<Person>) -> PyResult<HashMap<String, i32>> {
    Ok(v.into_iter().map(|p| (p.name, p.age)).collect())
}

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

ここでのポイントはもちろん#[derive(FromPyObject)]のついたPerson構造体です。
構造体のフィールドに対して#[pyo3(attribute("属性名"))]と付けることにより、入力された値から指定した属性の値を取ることができるので、nameとageの値を取って構造体にセットしているわけです。

このコードをビルドすることで、以下のようにPythonから呼び出すことが出来ます。

run.py
import pyo3_test

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

person_list = [Person("Alice", 20), Person("Bob", 30), Person("Charlie", 40)]
dic = pyo3_test.person_list_to_dict(person_list)
print(f'{dic}: {type(dic)}') # {'Bob': 30, 'Charlie': 40, 'Alice': 20}: <class 'dict'>

IntoPy: Rust から Python への型変換

続いて、IntoPyトレイトについてです。
自作のRustの型であるTomlValueをPythonに渡すことが出来れば本記事の目的[3]は達成したも同然なので、ここではそれに挑戦してみましょう。

なお、TomlValueというのは前回の記事で作った列挙型で、以下のように定義されています。

pub enum TomlValue {
    String(String),
    Integer(i64),
    Float(f64),
    Boolean(bool),
    OffsetDateTime(DateTime<chrono::offset::FixedOffset>),
    LocalDateTime(NaiveDateTime),
    LocalDate(NaiveDate),
    LocalTime(NaiveTime),
    Array(Vec<TomlValue>),
    InlineTable(HashMap<String, TomlValue>),
    Table(HashMap<String, TomlValue>),
    ArrayTable(Vec<HashMap<String, TomlValue>>),
}

TOMLというのは、パースした結果として連想配列に変換されるものと想定されている形式なのですが、その連想配列のvalueを表すのがこのTomlValueです。

このTomlValueに対してIntoPyトレイトを実装してみましょう。ユーザーガイドによると、IntoPy<PyObject>を実装するのが普通のようなので、それに従って実装してみます。

impl IntoPy<PyObject> for TomlValue {
    fn into_py(self, py: Python) -> PyObject {
        match self {
            TomlValue::String(s) => s.into_py(py),
            TomlValue::Integer(i) => i.into_py(py),
            TomlValue::Float(f) => f.into_py(py),
            TomlValue::Boolean(b) => b.into_py(py),
            TomlValue::OffsetDateTime(dt) => dt.into_py(py),
            TomlValue::LocalDateTime(dt) => dt.into_py(py),
            TomlValue::LocalDate(d) => d.into_py(py),
            TomlValue::LocalTime(t) => t.into_py(py),
            TomlValue::Array(a) => a.into_py(py),
            TomlValue::InlineTable(t) => t.into_py(py),
            TomlValue::Table(t) => t.into_py(py),
            TomlValue::ArrayTable(t) => t.into_py(py),
        }
    }
}

とても簡単ですね。TomlValueがラップしている各々の型がIntoPyを実装しているので、それらを取り出して単にinto_pyを呼び出すだけで済んでしまいました。

なお、注意点としては、日時を表すためにchronoの型 (DateTime<chrono::offset::FixedOffset>, NaiveDateTime, NaiveDate, NaiveTime) を使っているのですが、これらの型に対するIntoPy<PyObject>の実装はデフォルトでは存在しません。

chronoの型に対するIntoPy<PyObject>の実装を使えるようにするには、Cargo.tomlを修正してpyo3のfeaturesにchronoを追加する必要があります。

Cargo.toml
[package]
name = "pyo3_test"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "pyo3_test"
crate-type = ["cdylib"]

[dependencies]
-pyo3 = "0.18.3"
+pyo3 = { version = "0.18.3", features = ["chrono"] }
+chrono = "0.4"

余談: chronoの型に対しIntoPy<PyObject>を自分で実装する

余談になりますが、最初はこの features = ["chrono"]の存在を知らなくて、chronoの型に対するIntoPy<PyObject>の実装を自分で書いていました。
今回のケースだともうその実装を使う意味が無くなってしまったのですが、個人的には勉強になったのでそちらについても軽く紹介します。

その時の実装は以下の通りです。

impl IntoPy<PyObject> for TomlValue {
    fn into_py(self, py: Python) -> PyObject {
        match self {
            TomlValue::String(s) => s.into_py(py),
            TomlValue::Integer(i) => i.into_py(py),
            TomlValue::Float(f) => f.into_py(py),
            TomlValue::Boolean(b) => b.into_py(py),
            TomlValue::OffsetDateTime(dt) => {
                let locals = PyDict::new(py);
                let code = format!(
                    "import datetime; result = datetime.datetime({}, {}, {}, {}, {}, {}, {}, datetime.timezone(datetime.timedelta(hours={})))",
                    dt.year(),
                    dt.month(),
                    dt.day(),
                    dt.hour(),
                    dt.minute(),
                    dt.second(),
                    dt.nanosecond() / 1000,
                    dt.offset().local_minus_utc() / 3600,
                );
                locals.set_item("result", py.None()).unwrap();
                py.run(&code, None, Some(locals)).unwrap();
                let datetime = locals.get_item("result").unwrap();
                datetime.into()
            },
            TomlValue::LocalDateTime(dt) => {
                PyDateTime::new(
                    py,
                    dt.year(),
                    dt.month() as u8,
                    dt.day() as u8,
                    dt.hour() as u8,
                    dt.minute() as u8,
                    dt.second() as u8,
                    dt.nanosecond() / 1000,
                    None,
                )
                .unwrap()
                .into()
            },
            TomlValue::LocalDate(d) => {
                PyDate::new(
                    py,
                    d.year(),
                    d.month() as u8,
                    d.day() as u8
                )
                .unwrap()
                .into()
            },
            TomlValue::LocalTime(t) => {
                PyTime::new(
                    py,
                    t.hour() as u8,
                    t.minute() as u8,
                    t.second() as u8,
                    t.nanosecond() / 1000,
                    None,
                )
                .unwrap()
                .into()
            },
            TomlValue::Array(a) => a.into_py(py),
            TomlValue::InlineTable(t) => t.into_py(py),
            TomlValue::Table(t) => t.into_py(py),
            TomlValue::ArrayTable(t) => t.into_py(py),
        }
    }
}

特に良い経験になったと感じたのはOffsetDateTimeに対する実装の部分です。

TomlValue::OffsetDateTime(dt) => {
    let locals = PyDict::new(py);
    let code = format!(
        "import datetime; result = datetime.datetime({}, {}, {}, {}, {}, {}, {}, datetime.timezone(datetime.timedelta(hours={})))",
        dt.year(),
        dt.month(),
        dt.day(),
        dt.hour(),
        dt.minute(),
        dt.second(),
        dt.nanosecond() / 1000,
        dt.offset().local_minus_utc() / 3600,
    );
    locals.set_item("result", py.None()).unwrap();
    py.run(&code, None, Some(locals)).unwrap();
    let datetime = locals.get_item("result").unwrap();
    datetime.into()
},

+09:00みたいな形式のオフセットが付いたdatetimeをPyO3側でどうやって作るのかがちょっと調べてもよく分からなかったので、Python側でdatetimeオブジェクトを作って取り出して使っています。
into_pyの引数であるpy: Pythonという変数の.runというメソッドを使うことでPythonのコードを実行し、実行結果をローカル変数(locals)からget_itemで取り出して、それをintoでPyObjectに変換して返しています。

この実装は正直に言うとあまり良い実装ではない[4]と感じていますが、困った時にはこういう手段もあると分かったことでPyO3を使うことに対する不安は少し薄まった気がします。

エラーハンドリング

最後に、Rustのコード内で発生したエラーをPython側でハンドリングする方法について考えます。

ここまで、Rustで実装したPythonの関数はPyResult<T>を返すように実装してきました。
このPyResult<T>はResultの型エイリアスで、以下のように定義されています。

pub type PyResult<T> = Result<T, PyErr>;

この定義の中にあるPyErrはPythonの例外に対応しています。
つまり、エラー時に発生させるPythonの例外をRust側からコントロール出来るということです。

Rust側からPythonの例外を発生させる方法として、以下の3つのパターンについて考えてみます。

  1. Pythonの組み込みの例外を使いたい場合
  2. 自作の例外を使いたい場合
  3. std::from::From<E> for PyErr が既に実装されている場合

Pythonの組み込みの例外を使いたい場合

Rust側からPythonの組み込みの例外を発生させたい場合、PyO3には組み込みの例外に対応する型がそれぞれ用意されているので、単純にそれを使うだけでOKです。

以下の例では、ファイルのパスを受け取ってそのファイルを読み込む関数を実装しています。

lib.rs
use pyo3::exceptions::PyFileNotFoundError;
use pyo3::prelude::*;

#[pyfunction]
fn read_toml(toml_path: &str) -> PyResult<String> {
    match std::fs::read_to_string(toml_path) {
        Ok(toml) => Ok(toml),
        Err(e) => match e.kind() {
            std::io::ErrorKind::NotFound => Err(PyFileNotFoundError::new_err(format!(
                "File not found: {}",
                toml_path
            ))),
            _ => panic!("Unexpected error: {:?}", e),
        },
    }
}

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

ファイルが存在しない場合はPython側にFileNotFoundErrorが伝播するようにしたいので、それに対応するPyFileNotFoundErrorを返しています。
(それ以外のケースでは、どう処理すべきか決めていないのでとりあえずpanicするようにしています。)

実際にファイルが存在しない場合にこの関数を実行するとどうなるか見てみましょう。

run.py
import pyo3_test

try:
    # 存在しないファイルを指定
    print(pyo3_test.read_toml('not_exist.toml'))
except FileNotFoundError as e:
    print(f'Catched: {e}')

出力結果は以下の通りです。ちゃんと想定通りの挙動になっていることが分かります。

Catched: File not found: not_exist.toml

自作の例外を使いたい場合

組み込みの例外を使うのではなく、自作した例外を使いたいという場合もあると思います。
その場合は、以下のようにcreate_exception!マクロを使って例外を作成すればOKです。

lib.rs
use pyo3::create_exception;
use pyo3::prelude::*;

create_exception!(
    pyo3_test,
    TomlFileNotFoundError,
    pyo3::exceptions::PyException
);

#[pyfunction]
fn read_toml(toml_path: &str) -> PyResult<String> {
    match std::fs::read_to_string(toml_path) {
        Ok(toml) => Ok(toml),
        Err(e) => match e.kind() {
            std::io::ErrorKind::NotFound => Err(TomlFileNotFoundError::new_err(format!(
                "TOML file not found: {}",
                toml_path
            ))),
            _ => panic!("Unexpected error: {:?}", e),
        },
    }
}

/// A Python module implemented in Rust.
#[pymodule]
fn pyo3_test(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(read_toml, m)?)?;
    m.add(
        "TomlFileNotFoundError",
        _py.get_type::<TomlFileNotFoundError>(),
    )?;
    Ok(())
}

この例では、create_exception!マクロを使ってTomlFileNotFoundErrorという例外を作成しています。
例外を作ってしまえば、あとは組み込みの例外を発生させる時と同じようにnew_errで例外を作成して返すだけです。
Python内で例外をキャッチするために、モジュールに対してTomlFileNotFoundErrorを追加していることに注意してください。

Pythonから使う場合は、以下のようにすればOKです。先ほどの例のFileNotFoundErrorがpyo3_test.TomlFileNotFoundErrorに置き換わっているところだけが変更点です。

run.py
import pyo3_test

try:
    # 存在しないファイルを指定
    print(pyo3_test.read_toml('not_exist.toml'))
except pyo3_test.TomlFileNotFoundError as e:
    print(f'Catched: {e}')

実行結果は以下の通りです。

Catched: TOML file not found: not_exist.toml

std::from::From<E> for PyErr が既に実装されている場合

ここまではわざわざ自分でエラーを返してきたのですが、ユーザーガイドにも書かれている通り標準ライブラリのエラー型の多くにはstd::from::From<E> for PyErrが既に実装されています。
なので上の例であれば、例えば?を使って以下のように簡単にエラーを伝播させることが出来ます。

lib.rs
use pyo3::prelude::*;

#[pyfunction]
fn read_toml(toml_path: &str) -> PyResult<String> {
    let s = std::fs::read_to_string(toml_path)?;
    Ok(s)
}

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

Pythonからの呼び出し方もこれまでと同様です。

run.py
import pyo3_test

try:
    # 存在しないファイルを指定
    print(pyo3_test.read_toml('not_exist.toml'))
except FileNotFoundError as e:
    print(f'Catched: {e}')

実行結果は以下の通りです。

Catched: 指定されたファイルが見つかりません。 (os error 2)

Rustで開発したTOMLパーサーをPythonから呼び出す

ではいよいよTOMLパーサーをPythonから呼び出してみましょう。

本プロジェクトに含まれるファイルは以下の4つです。[5]

  • Cargo.toml
  • pyproject.toml
  • src/_toml_parser.rs
  • src/lib.rs

以下、それぞれ簡単に見ていきます。

Cargo.tomlとpyproject.toml

Cargo.toml と pyproject.toml は、上で説明に使ってきたものとほぼ同じです。[6]

Cargo.tomlとpyproject.toml
Cargo.toml
[package]
name = "toml_parser"
version = "0.1.0"
edition = "2021"

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

[dependencies]
pyo3 = { version = "0.18.3", features = ["chrono"] }
nom = "7.1.3"
chrono = "0.4"
maplit = "1.0"
pyproject.toml
[build-system]
requires = ["maturin>=0.15,<0.16"]
build-backend = "maturin"

[project]
name = "toml_parser"
requires-python = ">=3.7"
classifiers = [
    "Programming Language :: Rust",
    "Programming Language :: Python :: Implementation :: CPython",
    "Programming Language :: Python :: Implementation :: PyPy",
]


[tool.maturin]
features = ["pyo3/extension-module"]

src/_toml_parser.rs

src/_toml_parser.rs[7]には前回記事のTOMLパーサーの実装がそのまま書かれています。以下では、この中にあるparse_tomlという関数を使っていきます。

parse_tomlのシグネチャ
pub fn parse_toml(input: &str) -> MyResult<&str, HashMap<String, TomlValue>>

引数のinputはTOMLファイルの内容を表す文字列で、返り値のMyResultは標準のResultの型エイリアスです。
パースに成功した場合はHashMap<String, TomlValue>という連想配列を返し、失敗した場合はエラー情報を返します。[8]

src/lib.rs

src/lib.rsは、src/_toml_parser.rsのTOMLパーサーをPythonから呼び出すためのコードです。
Python側からはtoml_parserというモジュール名でimport出来るようになっています。

Python側に公開しているものは以下の3つです。

  • TomlParserError: TOMLパーサーでエラーが発生した場合に返される例外
  • parse_str: TOML形式の文字列をパースする関数
  • parse: TOMLファイルのパスを受け取って、そのファイルをパースする関数
src/lib.rs
mod _toml_parser;

use _toml_parser::{parse_toml, TomlValue};
use pyo3::create_exception;
use pyo3::prelude::*;
use std::collections::HashMap;

impl IntoPy<PyObject> for TomlValue {
    fn into_py(self, py: Python) -> PyObject {
        match self {
            TomlValue::String(s) => s.into_py(py),
            TomlValue::Integer(i) => i.into_py(py),
            TomlValue::Float(f) => f.into_py(py),
            TomlValue::Boolean(b) => b.into_py(py),
            TomlValue::OffsetDateTime(dt) => dt.into_py(py),
            TomlValue::LocalDateTime(dt) => dt.into_py(py),
            TomlValue::LocalDate(d) => d.into_py(py),
            TomlValue::LocalTime(t) => t.into_py(py),
            TomlValue::Array(a) => a.into_py(py),
            TomlValue::InlineTable(t) => t.into_py(py),
            TomlValue::Table(t) => t.into_py(py),
            TomlValue::ArrayTable(t) => t.into_py(py),
        }
    }
}

create_exception!(
    pyo3_toml_parser,
    TomlParserError,
    pyo3::exceptions::PyException
);

#[pyfunction]
fn parse_str(toml: &str) -> PyResult<HashMap<String, TomlValue>> {
    match parse_toml(toml) {
        Ok((_, map)) => Ok(map),
        Err(e) => Err(TomlParserError::new_err(e.to_string())),
    }
}

#[pyfunction]
fn parse(toml_path: &str) -> PyResult<HashMap<String, TomlValue>> {
    let toml = std::fs::read_to_string(toml_path)?;
    parse_str(&toml)
}

#[pymodule]
fn toml_parser(py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(parse_str, m)?)?;
    m.add_function(wrap_pyfunction!(parse, m)?)?;
    m.add("TomlParserError", py.get_type::<TomlParserError>())?;
    Ok(())
}

動作確認

テストケース的なものを幾つか用意して、簡単な動作確認をしてみます。

動作確認に使ったコード
check.py
import json
import toml_parser

print('Case1: Cargo.toml')
print(json.dumps(toml_parser.parse('Cargo.toml'), indent=4))

print('Case2: pyproject.toml')
print(json.dumps(toml_parser.parse('pyproject.toml'), indent=4))

print('Case3: datetime values')
toml_str = """
[オフセット付き日時]
odt1 = 1979-05-27T07:32:00Z
odt2 = 1979-05-27T00:32:00-07:00
odt3 = 1979-05-27T00:32:00.999999-07:00
odt4 = 1979-05-27 07:32:00Z

[ローカルの日時]
ldt1 = 1979-05-27T07:32:00
ldt2 = 1979-05-27T00:32:00.999999

[ローカルの日付]
ld1 = 1979-05-27

[ローカルの時刻]
lt1 = 07:32:00
lt2 = 00:32:00.999999
"""
d = toml_parser.parse_str(toml_str)

def convert_datetime(d):
    """
    datetime系の型をstrに変換 (そのままだとJSONに出来ない)
    """
    ret = {}
    for k, v in d.items():
        if isinstance(v, dict):
            ret[k] = convert_datetime(v)
        else:
            ret[k] = repr(v)
    return ret

d = convert_datetime(d)
print(json.dumps(d, indent=4, ensure_ascii=False))

print('Case4: Nested array-table')
toml_str = """
[[fruit]]
  name = "apple"

  [fruit.physical] # テーブル
    color = "red"
    shape = "round"

  [[fruit.variety]] # ネストされたテーブルの配列
    name = "red delicious"

  [[fruit.variety]]
    name = "granny smith"

[[fruit]]
  name = "banana"

  [[fruit.variety]]
    name = "plantain"
"""
print(json.dumps(toml_parser.parse_str(toml_str), indent=4))

print('Case5: DuplicationError')
toml_str = """
[a]
b = 1

[a]
b = 2
"""
try:
    toml_parser.parse_str(toml_str)
except toml_parser.TomlParserError as e:
    print(e)

print('Case6: InvalidTomlError')
toml_str = """
[a]
b = 1

[b
c = 2
"""
try:
    toml_parser.parse_str(toml_str)
except toml_parser.TomlParserError as e:
    print(e)

各ケースが雑に列挙されていて読みにくいと思いますので、1つ1つ見ていきましょう。

Case1: Cargo.toml

では試しに、このプロジェクトのCargo.tomlをパースしてみます。

print(json.dumps(toml_parser.parse('Cargo.toml'), indent=4))

なお、結果の表示にjson.dumpsを使っているのは、これを使うことで連想配列の構造を分かりやすく表示出来ると思ったからで、あまり深い意味はありません。[9]

実行結果は以下の通りです。ちゃんとCargo.tomlが表したかった連想配列の内容になっていますね。

実行結果
{
    "package": {
        "edition": "2021",
        "name": "toml_parser",
        "version": "0.1.0"
    },
    "lib": {
        "crate-type": [
            "cdylib"
        ],
        "name": "toml_parser"
    },
    "dependencies": {
        "pyo3": {
            "features": [
                "chrono"
            ],
            "version": "0.18.3"
        },
        "maplit": "1.0",
        "chrono": "0.4",
        "nom": "7.1.3"
    }
}

Case2: pyproject.toml

次に、このプロジェクトのpyproject.tomlをパースしてみます。

print(json.dumps(toml_parser.parse('pyproject.toml'), indent=4))

実行結果を見た感じ、こちらも大丈夫そうです。

実行結果
{
    "project": {
        "classifiers": [
            "Programming Language :: Rust",
            "Programming Language :: Python :: Implementation :: CPython",
            "Programming Language :: Python :: Implementation :: PyPy"
        ],
        "name": "toml_parser",
        "requires-python": ">=3.7"
    },
    "tool": {
        "maturin": {
            "features": [
                "pyo3/extension-module"
            ]
        }
    },
    "build-system": {
        "requires": [
            "maturin>=0.15,<0.16"
        ],
        "build-backend": "maturin"
    }
}

Case3: datetime values

次に、PyO3で実装する際に少し苦労した日時関連のvalueを含むTOMLをパースしてみましょう。

toml_str = """
[オフセット付き日時]
odt1 = 1979-05-27T07:32:00Z
odt2 = 1979-05-27T00:32:00-07:00
odt3 = 1979-05-27T00:32:00.999999-07:00
odt4 = 1979-05-27 07:32:00Z

[ローカルの日時]
ldt1 = 1979-05-27T07:32:00
ldt2 = 1979-05-27T00:32:00.999999

[ローカルの日付]
ld1 = 1979-05-27

[ローカルの時刻]
lt1 = 07:32:00
lt2 = 00:32:00.999999
"""
d = toml_parser.parse_str(toml_str)

def convert_datetime(d):
    """
    datetime系の型をstrに変換 (そのままだとJSONに出来ないので)
    """
    ret = {}
    for k, v in d.items():
        if isinstance(v, dict):
            ret[k] = convert_datetime(v)
        else:
            ret[k] = repr(v)
    return ret

d = convert_datetime(d)
print(json.dumps(d, indent=4, ensure_ascii=False))

補足: convert_datetimeという関数は、日時系の型を含むデータをjson.dumpsに渡すとエラーになるので、それを回避するために作ったものです。

実行結果を見てみると、ちゃんと妥当な結果になっているらしいことが分かります。(順番が色々と入れ替わっていてちょっと見づらいですが。)

実行結果
{
    "オフセット付き日時": {
        "odt4": "datetime.datetime(1979, 5, 27, 7, 32, tzinfo=datetime.timezone.utc)",
        "odt3": "datetime.datetime(1979, 5, 27, 7, 32, 0, 999999, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))",
        "odt2": "datetime.datetime(1979, 5, 27, 7, 32, tzinfo=datetime.timezone(datetime.timedelta(days=-1, seconds=61200)))",
        "odt1": "datetime.datetime(1979, 5, 27, 7, 32, tzinfo=datetime.timezone.utc)"
    },
    "ローカルの日時": {
        "ldt2": "datetime.datetime(1979, 5, 27, 0, 32, 0, 999999)",
        "ldt1": "datetime.datetime(1979, 5, 27, 7, 32)"
    },
    "ローカルの日付": {
        "ld1": "datetime.date(1979, 5, 27)"
    },
    "ローカルの時刻": {
        "lt1": "datetime.time(7, 32)",
        "lt2": "datetime.time(0, 32, 0, 999999)"
    }
}

Case4: Nested array-table

では、正常系の最後の例として、「ネストされたテーブルの配列」を含むTOMLをパースしてみましょう。[10]

toml_str = """
[[fruit]]
  name = "apple"

  [fruit.physical] # テーブル
    color = "red"
    shape = "round"

  [[fruit.variety]] # ネストされたテーブルの配列
    name = "red delicious"

  [[fruit.variety]]
    name = "granny smith"

[[fruit]]
  name = "banana"

  [[fruit.variety]]
    name = "plantain"
"""
print(json.dumps(toml_parser.parse_str(toml_str), indent=4))

この場合も、ちゃんと想定通りの結果になっています。

実行結果
{
    "fruit": [
        {
            "variety": [
                {
                    "name": "red delicious"
                },
                {
                    "name": "granny smith"
                }
            ],
            "physical": {
                "shape": "round",
                "color": "red"
            },
            "name": "apple"
        },
        {
            "name": "banana",
            "variety": [
                {
                    "name": "plantain"
                }
            ]
        }
    ]
}

Case5: DuplicationError

次に、エラーが発生する場合について見てみましょう。
以下の例は、キーに重複があるTOMLをパースする例です。

toml_str = """
[a]
b = 1

[a]
b = 2
"""
try:
    toml_parser.parse_str(toml_str)
except toml_parser.TomlParserError as e:
    print(e)

見てわかる通り、[a]というテーブルが2度書かれているのでこれはエラーにしないといけません。
自作のTOMLパーサーでは、このような場合にDuplicationErrorというエラーを返すようにしています。

実行結果は以下の通りです。aというキーが重複していることが分かるエラーが返ってきています。

実行結果
Parsing Failure: DuplicationError("key a is already defined")

Case6: InvalidTomlError

もう1つのエラー例として、形式的に不正なTOMLが与えられた場合について見てみましょう。
以下の例では、aとbという2つのテーブルが定義されていますが、bの方は括弧が閉じられておらず不正な形式になっています。

toml_str = """
[a]
b = 1

[b
c = 2
"""
try:
    toml_parser.parse_str(toml_str)
except toml_parser.TomlParserError as e:
    print(e)

実行結果は以下の通りです。ちょっと不親切ですが、bの方のテーブルのパースでエラーになったことが分かるメッセージになっています。

実行結果
Parsing Failure: InvalidTomlError("[b\nc = 2\n")

感想

やりたいことがそこまで難しくなかったこともあり、かなり簡単にPythonモジュール化が出来てしまいました。
もっと複雑なことをやろうとするとどうなるか分かりませんが、これだけ簡単に出来るとRustで実装したPythonのパッケージが色々と作られるのも分かる気がします。

脚注

  1. PyO3のREADMEにも具体例がたくさん書かれています: https://github.com/PyO3/pyo3#examples ↩︎

  2. 正確に言うと他にも幾つかファイルが作成されますが、主要なものはこの3つです。 ↩︎

  3. 自作TOMLパーサーの返り値であるHashMap<String, TomlValue>をRustからPythonに渡すこと。 ↩︎

  4. 呼び出す側のPython環境に依存する処理になってしまったりとか、コンパイラやrust-analyzerなどのサポートを受けられないとか、色々と悪い点がありそうです。 ↩︎

  5. リポジトリ: https://github.com/skwbc/toml_parser ↩︎

  6. 違う点は、package, lib, project の name が "toml_parser" になっているということと、前回記事で使っていたmaplitが依存関係に入っていることです。 ↩︎

  7. toml_parser.rsではなく_toml_parser.rsとアンダースコアが先頭に付いているのは、Python向けに作成するモジュールと名前が被ってしまってエラーになるからです。先頭にアンダースコアを付けるとPython感がかなり高まってしまいますが、こういう時Rustではどのようにすることが多いんでしょうね。 ↩︎

  8. この説明は正確ではないのですが、正確に説明しようとすると少し長くなるのでこうしました。興味がある人は前回記事を読んでみて下さい。 ↩︎

  9. printを使うとどういう構造になっているのかがぱっと見で全然分からないですし、pprintも試した感じちょっと微妙だったのです。 ↩︎

  10. TOMLパーサーを実装する上で最も厄介だったテストケースがこれなので、少し思い入れがあります。 ↩︎

Discussion