pyo3でRust側のエラーをPython側でキャッチする

今回は，pyo3クレートを用いてPythonからRustで定義した処理を呼び出すときのエラーハンドリングをしてみます．具体的にはRust側でthiserrorを用いて定義したエラーをPython側でキャッチします．

簡単なものですが，プログラムはこちらにあります．

エラーの定義

エラーは以下で定義しています．
列挙体はRust側でエラーハンドリングをするために作成しているので，Python側でのみエラーを扱うなら必要ありません．

error.rs

#[derive(thiserror::Error, Debug)]
pub enum MathError {
    #[error(transparent)]
    LogSmallError(#[from] LogSmallError),

    #[error(transparent)]
    ExpLargeError(#[from] ExpLargeError)
}

#[derive(thiserror::Error, Debug)]
#[error("too small x for log(x) x:{0}. (expected x > 0.01)")]
pub struct LogSmallError(pub f64);

#[derive(thiserror::Error, Debug)]
#[error("too large x for exp(x) x:{0}. (expected x < 10)")]
pub struct ExpLargeError(pub f64);

#[derive(thiserror::Error, Debug)]
#[error("too small abs y for x / y y:{0}. (expected y > 0.001)")]
pub struct DivSmallAbsError(pub f64);

RustのエラーをPythonからimportする

pyo3のユーザーガイドにあるように，まずcreate_exception!マクロでPythonから利用できるExceptionを定義します．第一引数はこのlib.rsで定義するモジュール名，第二引数はException名，第三引数は継承するスーパークラスです．ここでは簡単のためにerror.rsで定義したエラーと同じ名前としているので，lib.rsからerror.rsの同名のエラーを利用する場合はcrate::error::MathErrorのようにします．DivSmallAbsErrorについてはValueErrorとしてraiseさせるので，ここでは定義しません．

lib.rs

use pyo3::create_exception;
create_exception!(error_handling, MathError, pyo3::exceptions::PyException);

create_exception!(error_handling, LogSmallError, MathError);
create_exception!(error_handling, ExpLargeError, MathError);

次に各エラーからPyErrに変換できるようにFromトレイトを実装します．Fromトレイトを実装することで，pyo3::PyResult<T>（Rsult<T, PyErr>のエイリアス）に?で返せるようになります．DivSmallAbsErrorはValueErrorとしてraiseするためにpyo3::exceptions::PyValueErrorから作成しています．

lib.rs

impl From<crate::error::MathError> for PyErr {
    fn from(err: crate::error::MathError) -> PyErr {
        MathError::new_err(err.to_string())
    }
}

impl From<crate::error::LogSmallError> for PyErr {
    fn from(err: crate::error::LogSmallError) -> PyErr {
        LogSmallError::new_err(err.to_string())
    }
}


impl From<crate::error::ExpLargeError> for PyErr {
    fn from(err: crate::error::ExpLargeError) -> PyErr {
        ExpLargeError::new_err(err.to_string())
    }
}

impl From<crate::error::DivSmallAbsError> for PyErr {
    fn from(err: crate::error::DivSmallAbsError) -> PyErr {
        pyo3::exceptions::PyValueError::new_err(err.to_string())
    }
}

最後に，定義したExceptionをPythonから呼べるように登録します．これでエラーをExceptionとしてPythonからimportできるようになります．

lib.rs

#[pymodule]
fn error_handling(py: Python, m:&PyModule) -> PyResult<()> {
    m.add("MathError", py.get_type::<MathError>())?;
    m.add("LogSmallError", py.get_type::<LogSmallError>())?;
    m.add("ExpLargeError", py.get_type::<ExpLargeError>())?;


    ...その他関数の登録
}

Rust側でエラーを返しPythonでキャッチする

Fromトレイトを実装しているのでエラーを?で返せます．PyResult<T>の代わりにResult<T, crate::error::MathError>を返り値の型とすることもできますが，Python側でMathErrorとしてraiseされてしまいます．

lib.rs

fn log(x: f64) -> Result<f64, crate::error::LogSmallError> {
    if x < 0.01 {
        return Err(crate::error::LogSmallError(x));
    } else {
        return Ok(x.ln());
    }
}

fn exp(x: f64) -> Result<f64, crate::error::ExpLargeError> {
    if x > 10.0 {
        return Err(crate::error::ExpLargeError(x));
    } else {
        return Ok(x.exp());
    }
}

#[pyfunction]
fn rust_log(x: f64) -> PyResult<f64> {
    Ok(log(x)?)
}

#[pyfunction]
fn rust_exp(x: f64) -> PyResult<f64> {
    Ok(exp(x)?)
}

以上の関数をPython(jupyter)から呼んで例外を起こすと以下のようになります．

example.ipynb

from error_handling import MathError, LogSmallError, ExpLargeError
from error_handling import rust_log, rust_exp, rust_div, rust_log_with_exp

rust_log(0.0)

---------------------------------------------------------------------------

LogSmallError                             Traceback (most recent call last)

C:\Users\-\AppData\Local\Temp/ipykernel_10728/2422354142.py in <module>
----> 1 rust_log(0.0)


LogSmallError: too small x for log(x) x:0. (expected x > 0.01)

example.ipynb

try:
    rust_exp(1000)
except ExpLargeError as err:
    print(err)

too large x for exp(x) x:1000. (expected x < 10)

もちろん，スーパークラスを指定してもサブクラスをキャッチできます．

example.ipynb

try:
    rust_exp(1000)
except MathError as err:
    print(err)

too large x for exp(x) x:1000. (expected x < 10)

Fromトレイトの実装でPyValueErrorから作成したDivSmallAbsErrorはPythonではValueErrorがraiseされます．

lib.rs

fn div(x: f64, y: f64) -> Result<f64, crate::error::DivSmallAbsError> {
    if y < 0.001 {
        return Err(crate::error::DivSmallAbsError(y));
    } else {
        return Ok(x / y);
    }
}

#[pyfunction]
fn rust_div(x: f64, y: f64) -> PyResult<f64> {
    Ok(div(x, y)?)
}

example.ipynb

rust_div(1.0, 0.0)

---------------------------------------------------------------------------

ValueError                                Traceback (most recent call last)

C:\Users\-\AppData\Local\Temp/ipykernel_10728/109446888.py in <module>
----> 1 rust_div(1.0, 0.0)


ValueError: too small abs y for x / y y:0. (expected y > 0.001)

以上のようにpyo3ではRust側でも，Pythonで利用できるエラーが定義できます．さらに
Python3.10ではmatch式が導入されたので，rustでエラーコードを返すような以下の処理を

lib.rs

fn log_with_exp(x: f64) -> Result<(f64, f64), crate::error::MathError> {
    let log_x = log(x)?;
    let exp_x = exp(x)?;
    Ok((log_x, exp_x))
}

fn rust_handle_error(x: f64) -> i64 {
    let res = log_with_exp(x);
    if let Err(error) = res {
        match error {
            crate::error::MathError::LogSmallError(_) => { return 1;},
            crate::error::MathError::ExpLargeError(_) => { return 2;}
        }
    } else {
        return 0;
    }
}

match式を用いて以下のように簡単に実装できます．

example.ipynb

def handle_error(x: float) -> int:
    try:
        rust_log(x)
        rust_exp(x)
    except MathError as err:
        match err:
            case LogSmallError():
                return 1
            case ExpLargeError():
                return 2
            case _:
                return 100
    else:
        return 0