uvとPyO3でPythonからRustを呼ぶ②（基本的なデータ型の値をやり取りする）

# はじめに

備忘用。
この記事では、PyO3を使った、PythonとRust間での基本的なデータ型の変換方法について述べます。
勉強し始めの初心者のため、何か指摘や気づきがあれば是非コメントお願いします。

一応、以下の記事からの続きとなります。言語、ライブラリのバージョンやその他環境の詳細はそちらをご参照ください。

- [uvとPyO3でPythonからRustを呼ぶ①（プロジェクト作成・Hello World実行まで）](https://qiita.com/nukipei/items/f096a1df6c8074b16150)

# 基本的なデータ型の変換

なお、その他の型については、PyO3のGitHub[^1]や公式ガイド[^2]、APIリファレンス[^3]をご参照ください。

## 数値

Rustの整数型と浮動小数点数型は対応するPythonの数値型に自動的に変換されます。
次に示すコードは、2つの数値を引数として受け取り、足し算を行う3つの関数を定義しています。
+ `add_int32_numbers`: 整数型（32ビット）の足し算を行う
+ `add_int64_numbers`: 整数型（64ビット）の足し算を行う
+ `add_float32_numbers`: 浮動小数点数型（32ビット）の足し算を行う

```rust: src/lib.rs
use pyo3::prelude::*;

#[pyfunction]
// 整数型（32ビット）の足し算を行う
fn add_int32_numbers(a: i32, b: i32) -> i32 {
    a + b
}

#[pyfunction]
// 整数型（64ビット）の足し算を行う
fn add_int64_numbers(a: i64, b: i64) -> i64 {
    a + b
}

#[pyfunction]
// 浮動小数点数型（32ビット）の足し算を行う
fn add_float32_numbers(a: f32, b: f32) -> f32 {
    a + b
}

#[pymodule]
fn _core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(add_int32_numbers, m)?)?;
    m.add_function(wrap_pyfunction!(add_int64_numbers, m)?)?;
    m.add_function(wrap_pyfunction!(add_float32_numbers, m)?)?;
    Ok(())
}
```

```python: src/example_ext/__init__.py
from example_ext._core import add_int32_numbers, add_int64_numbers, add_float32_numbers 

def main() -> None:
    print(add_int32_numbers(1, 2))
    print(add_int64_numbers(1, 2))
    print(add_float32_numbers(1.5, 1.5))
```

↑のPythonコードの実行結果は以下の通りです。
```
>uv run --reinstall --directory src/example_ext example-ext
3
3
3.0
```

注意点として`add_float32_numbers`の引数に整数を指定できますが、一方、`add_int32_numbers`と`add_int64_numbers`の引数には浮動小数点数型を指定できません。そのため、以下のようにfloat型を引数に指定するとTypeErrorが発生します。
```python
print(add_int32_numbers(1.0, 2))  # TypeError: argument 'a': 'float' object cannot be interpreted as an integer
print(add_int64_numbers(1.0, 2))  # TypeError: argument 'a': 'float' object cannot be interpreted as an integer
print(add_float32_numbers(1, 2))  # 3
```

各数値型の上下限値にも注意する必要があります。例えば、`add_int32_numbers`の引数に2147483647より大きい値を指定すると、OverflowErrorが発生します。これは、add_int32_numbersの（Rust側の）引数の型が`i32`であるためです。（実際、`i64`である`add_int64_numbers`は正常に実行されます。）

```python
print(add_int32_numbers(1, 2147483647))  # OverflowError: Python int too large to convert to C long
print(add_int64_numbers(1, 2147483648))  # 2147483649
```

また、以下のように、引数に上下限値内の値を指定したとしても、Rust側での計算の結果、オーバーフローとなることもあります。
```python
print(add_int32_numbers(1, 2147483647))  # -2147483648
print(add_int64_numbers(1, 2147483647))  # 2147483649
```


## 文字列

RustのString/&strはPythonのstrに変換できます。以下のコードは、引数nameと"Hello, {}!"を結合し文字列型として返します。

```rust: src/lib.rs
use pyo3::prelude::*;


#[pyfunction]
// 引数nameと"Hello, {}!"を結合し返す
fn greet(name: &str) -> String {
    format!("Hello, {}!", name)
}

#[pymodule]
fn _core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(greet, m)?)?;
    Ok(())
}
```

```python: src/example_ext/__init__.py
from example_ext._core import greet

def main() -> None:
    print(greet("World"))
    print(greet("Bob"))
```

実行結果は以下の通りです。
```
>uv run --reinstall --directory src/example_ext example-ext
Hello, World!
Hello, Bob!
```

## list

### 全ての要素の型が同じとき

Pythonのlistは、全ての要素の型が同じであり、かつRustに対応する型があれば、RustのVec<T>と相互変換できます。以下は整数を要素とするlist/Vecに対し、各要素を2倍にして返す関数`double_list`を定義しています。

```rust: src/lib.rs
use pyo3::prelude::*;


#[pyfunction]
// Vecの各要素を2倍にして返す
fn double_list(values: Vec<i32>) -> Vec<i32> {
    values.iter().map(|x| x * 2).collect()
}

#[pymodule]
fn _core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(double_list, m)?)?;
    Ok(())
}
```

```python: src/example_ext/__init__.py
from example_ext._core import double_list

def main() -> None:
    print(double_list([1, 2, 3])) 
```

実行結果は以下の通りです。
```
>uv run --reinstall --directory src/example_ext example-ext
[2, 4, 6]
```

### 要素の型が異なるとき

ところで、Pythonのlistは異なる型の要素を含むことができますが、`double_list`の引数としてそのようなlistを指定すると、当然TypeErrorが発生します。

```python
print(double_list([1, "2", 3]))   # TypeError: argument 'values': must be real number, not str
```

そこで、PyO3の`PyList`[^4]を活用することで、複数の型を含むPythonのlistもRust上で扱えるようになります。以下のコードはlistを受け取り、int型の要素のみを2倍して返します。

```rust: src/lib.rs
#[pyfunction]
// int型の要素のみを2倍する
fn double_list_only_int(py: Python, list: &Bound<'_, PyList>) -> PyResult<Py<PyList>> {
    // PyListを要素無しで定義
    let mut result = PyList::empty_bound(py);

    for value in list {
        // int型（PyInt型）の場合、2をかける
        if value.is_instance_of::<PyInt>() {
            result.append(value.extract::<i64>()? * 2)?;
        } 
        // それ以外の場合、そのまま。
        else {
            result.append(value)?;
        }
    }
    // resultはpyo3::Bound<'_, PyList>型なので、intoメソッドでpyo3::Py<PyList>型に変換
    Ok(result.into())
}

#[pymodule]
fn _core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(double_list_only_int, m)?)?;
    Ok(())
}
```

```python: src/example_ext/__init__.py
from example_ext._core import double_list

def main() -> None:
    print(double_list_only_int([1, "2", 3]))  # [2, '2', 6]
    print(double_list_only_int([1, "2", 3, [1, 2], ("1", 2), {"key": "value"}]))
    # [2, '2', 6, [1, 2], ('1', 2), {'key': 'value'}]
```

## dict

### 全てのキー、バリューの型がそれぞれ同じとき

Pythonのdictは、全てのキー、バリューの型がそれぞれ同じであり、かつそれらの型に対応する型がRustにあれば、HashMapと相互変換できます。以下のコードは、dictの各バリューに指定されたテキストの単語数をカウントし、キーが元のdictのキーで、バリューが単語数であるdictを返します。

```rust: src/lib.rs
use std::collections::HashMap;


#[pyfunction]
// dictの各バリューの単語数をカウントする
fn count_words(text_map: HashMap<String, String>) -> HashMap<String, i32> {
    let mut result: HashMap<String, i32> = HashMap::new();
    
    for (key, text) in text_map {
        let mut count = 0;
        for _word in text.split_whitespace() {
            count += 1;
        }
        result.insert(key, count);
    }
    
    result
}
```

```python: src/example_ext/__init__.py
from example_ext._core import count_words

def main() -> None:
    print(count_words({"text1": "I am Nukipei", "text2": "Hello World", "text3": "convert dict to HashMap"}))
```

実行結果は以下の通りです。
```
>uv run --reinstall --directory src/example_ext example-ext
{'text1': 3, 'text3': 4, 'text2': 2}
```

### キー、バリューの型がそれぞれ異なるとき

キー、バリューの型が異なる場合は、PyO3の`PyDict`[^5]を活用できます。
コードについては、省略しますが、使用感は先ほど紹介したPyListと似た感じかと思います。

## Optional

Noneを取り得る型である、PythonのOptional型はRustのOption型と対応します。以下のコードは、引数に指定したlist/Vecに対し、最初に見つけた偶数を返し、すべて奇数であれば、Noneを返します。

```rust: src/lib.rs
use pyo3::prelude::*;


#[pyfunction]
// 最初に見つけた偶数を返す。存在しない場合はNoneを返す
fn find_first_even(values: Vec<i32>) -> Option<i32> {
    values.into_iter().find(|x| x % 2 == 0)
}


#[pymodule]
fn _core(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(find_first_even, m)?)?;
    Ok(())
}
```


```python: src/example_ext/__init__.py
from example_ext._core import find_first_even

def main() -> None:
    print(find_first_even([1, 2, 3]))  # 偶数を含む
    print(find_first_even([1, 3, 5]))  # すべて奇数
```

実行結果は以下の通りです。すべて奇数の時はNoneを返すことが確認できます。
```
>uv run --reinstall --directory src/example_ext example-ext
2
None
```

ちなみに、`#[pyo3(signature = (...))]`[^6]を利用することで、デフォルト値を設定することができます。
以下の関数は、先述の`find_first_even`の修正版であり、デフォルト値がfalseであるbool型の引数reverseを追加しています。なお、reverseがfalseの時はVecの先頭から偶数を探し、reverseがtrueの時は末尾から偶数を探します。

```rust
#[pyfunction]
#[pyo3(signature = (values, reverse=false))]
// 最初に見つけた偶数を返す。存在しない場合はNoneを返す
fn find_first_even(values: Vec<i32>, reverse: bool) -> Option<i32> {
    if reverse {
        values.into_iter().rev().find(|x| x % 2 == 0)
    } else {
        values.into_iter().find(|x| x % 2 == 0)
    }
}
```

```python
print(find_first_even([1, 2, 3, 4]))  # 2
print(find_first_even([1, 2, 3, 4], reverse=False))  # 2
print(find_first_even([1, 2, 3, 4], reverse=True))  # 4
```

[^1]: [PyO3 GitHub](https://github.com/PyO3/pyo3)
[^2]: [PyO3公式ガイド](https://pyo3.rs/)
[^3]: [PyO3APIリファレンス](https://docs.rs/pyo3/)
[^4]: https://docs.rs/pyo3/0.22.4/pyo3/types/struct.PyList.html
[^5]: https://docs.rs/pyo3/0.22.4/pyo3/types/struct.PyDict.html
[^6]: https://pyo3.rs/v0.22.4/function/signature.html#using-pyo3signature--