PythonのコードをRustから使う その５ Pythonのファイルを読み込む

前回の記事

概要

RustからPythonのコードを呼ぶ時に、文字列を直接ベタ書きしてたのを、ただファイルから読み込むよ、というだけです。

シンプルな例

ファイル構造

$ tree
.
├── Cargo.toml
├── scripts
│   └── mymodule.py
└── src
    └── main.rs

scriptsの中にmymodule.pyを配置。

scripts/mymodule.py

def function1(arg0, arg1):
    print("function1() called. add two args.")
    return {
        "sub": arg0 - arg1,
        "add": arg0 + arg1,
        "str": f"arg0 is '{arg0}', arg1 is '{arg1}'",
        }

src/main.rs

use std::{collections::HashMap, fs, path::PathBuf};

use pyo3::{prelude::*, types::{PyDict, PyFloat, PyInt, PyList, PyNone, PyString, PyTuple}};
use serde_json::{json, Value}; // 主要なモジュール

fn main() {
    let mut rust_value: Value = json!({});
    let arg0 = 1;
    let arg1 = 2;

    // ロードするPythonファイルのフルパス
    let path_to_pyfile = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("scripts/mymodule.py");

    // Python::with_gilでPython環境呼び出す
    let pyresult  = Python::with_gil(|py| -> PyResult<Py<PyAny>> {
        // コードを文字列に変換
        let app_code = fs::read_to_string(path_to_pyfile)?;
        // ここでコードをモジュールとしてロード。ロードした時にグローバルに置いたprintの行が実行される
        let module = PyModule::from_code_bound(py, 
            app_code.as_str(), "", "")?;
        // ここでモジュールの中のfunction1にアクセスできるようになる
        let app_func: Py<PyAny> = module.getattr("function1")?.into();
        // 引数をPyTupleでまとめて呼び出し。返却値を受け取る。
        let func_result = app_func.call1(py, PyTuple::new_bound(py, [arg0, arg1]))?;

        let ref_pydict = func_result.clone().extract::<&PyDict>(py)?; // 結果をPyDictの参照に変換する
        rust_value = pydict_to_value(ref_pydict)?; // この関数でPyDictをserde_json::Valueに変換する。
        Ok(func_result)
    });
    println!("pyresult is {pyresult:?}");
    println!("rust_value is {rust_value:?}");
}

// 個々のPyAny型変数の型を確認しながら変換する
pub fn pyany_to_value(value: &PyAny) -> PyResult<Value> {
    if value.is_instance_of::<PyString>() {
        Ok(Value::from(value.extract::<String>()?))
    } else if value.is_instance_of::<PyFloat>() {
        Ok(Value::from(value.extract::<f64>()?))
    } else if value.is_instance_of::<PyInt>() {
        Ok(Value::from(value.extract::<i64>()?))
    } else if value.is_instance_of::<PyList>() {
        pylist_to_value(value.extract::<&PyList>()?)
    } else if value.is_instance_of::<PyDict>() {
        pydict_to_value(value.extract::<&PyDict>()?)
    } else if value.is_instance_of::<PyNone>() {
        Ok(Value::Null)
    } else {
        Err(pyo3::exceptions::PyException::new_err("未対応のタイプなので失敗"))
    }
}

// リストをValueに変換
fn pylist_to_value(pylist: &PyList) -> PyResult<Value> {
    let mut vec: Vec<Value> = Vec::new();
    for value in pylist.into_iter() {// 中身を一個ずつPyAnyからValueに変換する
        vec.push(pyany_to_value(value)?);
    }
    Ok(vec.into())
}

// PyDictをValueに変換
pub fn pydict_to_value(pydict: &PyDict) -> PyResult<Value> {
    let mut map: HashMap<String, Value> = HashMap::new();
    for (key, value) in pydict.into_iter() { // 中身を一個ずつPyAnyからValueに変換する
        map.insert(key.extract::<String>()?, pyany_to_value(value)?);
    }
    Ok(json!(map))
}

実行

function1() called. add two args.
pyresult is Ok(Py(0x100fa3b00))
rust_value is Object {"add": Number(3), "str": String("arg0 is '1', arg1 is '2'"), "sub": Number(-1)}

至ってシンプルだと思う。

すこし複雑な例

ここでPythonのコードをもうちょっと構造化したくなるのが人の心というもの。

ファイル構造

$ tree
.
├── Cargo.toml
├── scripts
│   ├── mymodule.py
│   └── othermodule
│       └── __init__.py
└── src
    └── main.rs

othermoduleといモジュールを作っておいた。

scripts/othermodule/init.py

def other_module_function():
    return "This is other_module_function result."

scripts/mymodule.py

import othermodule #この行を追加

def function1(arg0, arg1):
    print("function1() called. add two args.")
    return {
        "sub": arg0 - arg1,
        "add": arg0 + arg1,
        "str": f"arg0 is '{arg0}', arg1 is '{arg1}'",
        "othermodule": othermodule.other_module_function() # この行を追加
        }

src/main.rs

use std::{collections::HashMap, fs, path::PathBuf};

use pyo3::{prelude::*, types::{PyDict, PyFloat, PyInt, PyList, PyNone, PyString, PyTuple}};
use serde_json::{json, Value}; // 主要なモジュール

fn main() {
    let mut rust_value: Value = json!({});
    let arg0 = 1;
    let arg1 = 2;

    // ロードするPythonファイルのフルパス
    let path_to_pyfile = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("scripts/mymodule.py");
    // Python::with_gilでPython環境呼び出す
    let pyresult: Result<Py<PyAny>, PyErr>  = Python::with_gil(|py| -> PyResult<Py<PyAny>> {
        let parent = path_to_pyfile.parent().unwrap().to_str().unwrap(); // 指定されたPythonファイルがあるディレクトリのパスを取得
        // ここでsys.pathに取得したparentの値を追加すると、ここからimportできる。
        let _ = PyModule::from_code_bound(py, &format!(r#"
import sys
if not "{parent:}" in sys.path:
    sys.path.append("{parent:}")
            "#), "", "");
        
        // コードを文字列に変換
        let app_code = fs::read_to_string(path_to_pyfile)?;
        // ここでコードをモジュールとしてロード。ロードした時にグローバルに置いたprintの行が実行される
        let module = PyModule::from_code_bound(py, 
            app_code.as_str(), "", "")?;
        // ここでモジュールの中のfunction1にアクセスできるようになる
        let app_func: Py<PyAny> = module.getattr("function1")?.into();
        // 引数をPyTupleでまとめて呼び出し。返却値を受け取る。
        let func_result = app_func.call1(py, PyTuple::new_bound(py, [arg0, arg1]))?;

        let ref_pydict = func_result.clone().extract::<&PyDict>(py)?; // 結果をPyDictの参照に変換する
        rust_value = pydict_to_value(ref_pydict)?; // この関数でPyDictをserde_json::Valueに変換する。
        Ok(func_result)
    });
    println!("pyresult is {pyresult:?}");
    println!("rust_value is {rust_value:?}");
}

// 個々のPyAny型変数の型を確認しながら変換する
pub fn pyany_to_value(value: &PyAny) -> PyResult<Value> {
    if value.is_instance_of::<PyString>() {
        Ok(Value::from(value.extract::<String>()?))
    } else if value.is_instance_of::<PyFloat>() {
        Ok(Value::from(value.extract::<f64>()?))
    } else if value.is_instance_of::<PyInt>() {
        Ok(Value::from(value.extract::<i64>()?))
    } else if value.is_instance_of::<PyList>() {
        pylist_to_value(value.extract::<&PyList>()?)
    } else if value.is_instance_of::<PyDict>() {
        pydict_to_value(value.extract::<&PyDict>()?)
    } else if value.is_instance_of::<PyNone>() {
        Ok(Value::Null)
    } else {
        Err(pyo3::exceptions::PyException::new_err("未対応のタイプなので失敗"))
    }
}

// リストをValueに変換
fn pylist_to_value(pylist: &PyList) -> PyResult<Value> {
    let mut vec: Vec<Value> = Vec::new();
    for value in pylist.into_iter() {// 中身を一個ずつPyAnyからValueに変換する
        vec.push(pyany_to_value(value)?);
    }
    Ok(vec.into())
}

// PyDictをValueに変換
pub fn pydict_to_value(pydict: &PyDict) -> PyResult<Value> {
    let mut map: HashMap<String, Value> = HashMap::new();
    for (key, value) in pydict.into_iter() { // 中身を一個ずつPyAnyからValueに変換する
        map.insert(key.extract::<String>()?, pyany_to_value(value)?);
    }
    Ok(json!(map))
}

cargo run

$ cargo run
function1() called. add two args.
pyresult is Ok(Py(0x100c60cc0))
rust_value is Object {"add": Number(3), "othermodule": String("This is other_module_function result."), "str": String("arg0 is '1', arg1 is '2'"), "sub": Number(-1)}

import sysして、sys.path.appendを呼んで無理やりルートディレクトリ追加したら動いた。

動いた。うん。

次の記事