如何使用 pyo3 将 Rust 函数作为回调传递给 Python

Question

我使用 Pyo3 从 Python 调用 Rust 函数，反之亦然。

我正在努力实现以下目标：

• Python 调用rust_function_1

• Rust 函数rust_function_1调用 Python 函数python_function，传递 Rust 函数rust_function_2作为回调参数

• Python 函数python_function调用回调，在本例中是 Rust 函数rust_function_2

我无法弄清楚如何rust_function_2作为回调参数传递给python_function.

我有以下Python代码：

import rust_module

def python_function(callback):
    print("This is python_function")
    callback()

if __name__ == '__main__':
    rust_module.rust_function_1()

我有以下非编译 Rust 代码：

use pyo3::prelude::*;

#[pyfunction]
fn rust_function_1() -> PyResult<()> {
    println!("This is rust_function_1");
    Python::with_gil(|py| {
        let python_module = PyModule::import(py, "python_module")?;
        python_module
            .getattr("python_function")?
            .call1((rust_function_2.into_py(py),))?;  // Compile error
        Ok(())
    })
}

#[pyfunction]
fn rust_function_2() -> PyResult<()> {
    println!("This is rust_function_2");
    Ok(())
}

#[pymodule]
#[pyo3(name = "rust_module")]
fn quantum_network_stack(_python: Python, module: &PyModule) -> PyResult<()> {
    module.add_function(wrap_pyfunction!(rust_function_1, module)?)?;
    module.add_function(wrap_pyfunction!(rust_function_2, module)?)?;
    Ok(())
}

错误信息是：

error[E0599]: the method `into_py` exists for fn item `fn() -> Result<(), PyErr> {rust_function_2}`, but its trait bounds were not satisfied
  --> src/lib.rs:10:37
   |
10 |             .call1((rust_function_2.into_py(py),))?;
   |                                     ^^^^^^^ method cannot be called on `fn() -> Result<(), PyErr> {rust_function_2}` due to unsatisfied trait bounds
   |
   = note: `rust_function_2` is a function, perhaps you wish to call it
   = note: the following trait bounds were not satisfied:
           `fn() -> Result<(), PyErr> {rust_function_2}: AsPyPointer`
           which is required by `&fn() -> Result<(), PyErr> {rust_function_2}: pyo3::IntoPy<Py<PyAny>>`

Answer 1

PitaJ 的评论引导我找到了解决方案。

有效的 Rust 代码：

use pyo3::prelude::*;

#[pyclass]
struct Callback {
    #[allow(dead_code)] // callback_function is called from Python
    callback_function: fn() -> PyResult<()>,
}

#[pymethods]
impl Callback {
    fn __call__(&self) -> PyResult<()> {
        (self.callback_function)()
    }
}

#[pyfunction]
fn rust_function_1() -> PyResult<()> {
    println!("This is rust_function_1");
    Python::with_gil(|py| {
        let python_module = PyModule::import(py, "python_module")?;
        let callback = Box::new(Callback {
            callback_function: rust_function_2,
        });
        python_module
            .getattr("python_function")?
            .call1((callback.into_py(py),))?;
        Ok(())
    })
}

#[pyfunction]
fn rust_function_2() -> PyResult<()> {
    println!("This is rust_function_2");
    Ok(())
}

#[pymodule]
#[pyo3(name = "rust_module")]
fn quantum_network_stack(_python: Python, module: &PyModule) -> PyResult<()> {
    module.add_function(wrap_pyfunction!(rust_function_1, module)?)?;
    module.add_function(wrap_pyfunction!(rust_function_2, module)?)?;
    module.add_class::<Callback>()?;
    Ok(())
}

有效的Python代码（与问题中相同）：

import rust_module

def python_function(callback):
    print("This is python_function")
    callback()

if __name__ == '__main__':
    rust_module.rust_function_1()

以下解决方案以多种方式改进了上述解决方案：

• Rust 提供的内容callback会被存储并稍后调用，而不是立即调用（这对于现实生活中的用例来说更现实）

• 每次 Python 调用 Rust 时，它都会传入一个PythonApi对象，这样 Rust 函数就不需要在import每次调用时都执行 Python 操作了。

• Rust 提供的回调除了普通函数之外还可以是捕获变量（仅移动语义）的闭包。

更通用的Rust代码如下：

use pyo3::prelude::*;

#[pyclass]
struct Callback {
    #[allow(dead_code)] // callback_function is called from Python
    callback_function: Box<dyn Fn(&PyAny) -> PyResult<()> + Send>,
}

#[pymethods]
impl Callback {
    fn __call__(&self, python_api: &PyAny) -> PyResult<()> {
        (self.callback_function)(python_api)
    }
}

#[pyfunction]
fn rust_register_callback(python_api: &PyAny) -> PyResult<()> {
    println!("This is rust_register_callback");
    let message: String = "a captured variable".to_string();
    Python::with_gil(|py| {
        let callback = Box::new(Callback {
            callback_function: Box::new(move |python_api| {
                rust_callback(python_api, message.clone())
            }),
        });
        python_api
            .getattr("set_callback")?
            .call1((callback.into_py(py),))?;
        Ok(())
    })
}

#[pyfunction]
fn rust_callback(python_api: &PyAny, message: String) -> PyResult<()> {
    println!("This is rust_callback");
    println!("Message = {}", message);
    python_api.getattr("some_operation")?.call0()?;
    Ok(())
}

#[pymodule]
#[pyo3(name = "rust_module")]
fn quantum_network_stack(_python: Python, module: &PyModule) -> PyResult<()> {
    module.add_function(wrap_pyfunction!(rust_register_callback, module)?)?;
    module.add_function(wrap_pyfunction!(rust_callback, module)?)?;
    module.add_class::<Callback>()?;
    Ok(())
}

比较通用的Python代码如下：

import rust_module


class PythonApi:

    def __init__(self):
        self.callback = None

    def set_callback(self, callback):
        print("This is PythonApi::set_callback")
        self.callback = callback

    def call_callback(self):
        print("This is PythonApi::call_callback")
        assert self.callback is not None
        self.callback(self)

    def some_operation(self):
        print("This is PythonApi::some_operation")

def python_function(python_api, callback):
    print("This is python_function")
    python_api.callback = callback


def main():
    print("This is main")
    python_api = PythonApi()
    print("Calling rust_register_callback")
    rust_module.rust_register_callback(python_api)
    print("Returned from rust_register_callback; back in main")
    print("Calling callback")
    python_api.call_callback()


if __name__ == '__main__':
    main()

后一版本代码的输出如下：

This is main
Calling rust_register_callback
This is rust_register_callback
This is PythonApi::set_callback
Returned from rust_register_callback; back in main
Calling callback
This is PythonApi::call_callback
This is rust_callback
Message = a captured variable
This is PythonApi::some_operation