Search code examples
pythonrustpyo3

Calling a Rust function that calls a Python function from Python code


I’m new to pyo3, and I would like to achieve the following thing. My main code is in Python, but I would like to call a Rust function from this Python code. The trick is, the Rust function should be able to call a Python function provided as an argument.

Let’s say that I have a Python function, that cannot be embedded inside of the Rust module:

def double(x):
   return 2*x

I want to create a Rust function taking this function as an argument:

use pyo3::prelude::*;

fn my_rust_function(a: usize, f: ...) -> PyResult<usize> {
   ...
}

Finally, this Rust function would be used somewhere in the same Python code as above:

import rust_module

def double(x):
   return 2*x

def main():
   return rust_module.my_rust_function(3, double)

Would such a behavior be possible? I read parts of the documentation about how to call Rust from Python and Python from Rust, but never such a double-way use.


Solution

  • In my_rust_function(), take Bound<'_, PyAny> (or any variant, such as Borrowed<'_, '_, PyAny> or Py<PyAny>, but prefer Bound unless you have reasons not to), and use the call0() method on it to call it without arguments, or call1() to call it with positional arguments only, or call() to call it with both positional and keyword arguments.

    Exporting the Rust function to Python should be done as usual.

    For example:

    use pyo3::prelude::*;
    
    #[pyfunction]
    fn my_rust_function(a: usize, f: Bound<'_, PyAny>) -> PyResult<usize> {
        f.call1((a,))?.extract::<usize>()
    }
    
    #[pymodule]
    fn rust_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
        m.add_function(wrap_pyfunction!(my_rust_function, m)?)?;
        Ok(())
    }