Support for wrapping rust closures as python functions #1901

LaurentMazare · 2021-10-03T20:19:16Z

Hello and thanks for all the work on this very useful crate!
My understanding is that there is currently no easy way to wrap Rust closures and call them from Python. This PR adds a new function PyCFunction::new_closure to help doing so. Note that I'm not very familiar with the pyo3 internals, so it's unclear to me whether this is the right api for exposing this, any feedback is very welcome, and maybe this should be discussed in a 'need-design' issue first (though I haven't found any related such issue).

Technically, this works by boxing the Rust closure twice (to get around the fat pointer size) and the resulting pointer is stored in a Python capsule. Then PyCFunction_NewEx is used to create a Python callable function, this uses a 'dispatch' function run_closure that unpacks the capsule and calls the closure. The capsule also registers a destructor to collect the closure when it's not used anymore.
There are two small tests, one of them where the closure mutates some Rust based state.

Current limitations include:

The wrapped closure gets both the positional and keyword arguments, it would be possible to also have a function that only handles positional arguments or even no argument.
There is a RefUnwindSafe constraint that is used for handle_panic. I'm not sure if it's possible to remove it somehow, if this constraint is necessary it should probably be added to the new_closure function too.
The function name and documentation are hard-coded, maybe having a version of PyMethodDef using owned string instead of &'static str could help get around this and allow the user to specify the function name and documentation.
Again happy to get feedback on any of these or other limitations I have missed (or feel free to edit the PR too).

mejrs

Thanks 👍 I really want an api such as this. Thank you for doing this 🙏🏻

However there are a couple of issues with this implementation.

For starters, you can do use-after-frees by transferring references into the closure:

use pyo3::prelude::*;
use pyo3::types::{PyCFunction, PyDict, PyTuple};

fn main() {
    let fun = Python::with_gil(|py| {
        let stuff = vec![0, 1, 2, 3, 4];
        let ref_: &[u8] = &stuff;

        let counter_fn = |_args: &PyTuple, _kwargs: Option<&PyDict>| {
            println!("called");
            println!("This is five: {:?}", ref_.len());

            Ok("".into_py(py))
        };

        let counter_py: Py<PyCFunction> = PyCFunction::new_closure(counter_fn, py).unwrap().into();

        counter_py
    });

    Python::with_gil(|py| {
        fun.call0(py).unwrap();
    });
}

This is five: 738180135880

Similarly you can probably (I haven't checked) use Py<PyCFunction> to smuggle non-Send types to other threads. So you'll want (at least) a Send + 'static bound.

The Box<dyn Fn(&types::PyTuple, Option<&types::PyDict>) -> PyResult<PyObject>> bound is quite rigid. It would be nice to have closures that don't need args/kwargs arguments or need to return a PyResult. Having to return Ok(py.None()) is going to get annoying. Maybe we can do something clever with traits here?

In the past I've used a pyclass with a call method for this, like so:

use pyo3::prelude::*;
use pyo3::types::{PyDict, PyTuple};


type C = Box<dyn Fn(&PyTuple, Option<&PyDict>) -> PyResult<()> + Send + 'static>;

#[pyclass]
struct Closure {
    wraps: C,
}

#[pymethods]
impl Closure {
    #[call]
    #[args(args = "*", kwargs = "**")]
    fn __call__(
        &self,
        args: &PyTuple,
        kwargs: Option<&PyDict>,
    ) -> PyResult<()> {
        (self.wraps)(args, kwargs)
    }
}

fn main() {
    Python::with_gil(|py| {
        let f = |_: &PyTuple, _: Option<&PyDict>| {
            println!("called");
            Ok(())
        };

        let closure = Closure{wraps: Box::new(f)};
		closure.__call__(PyTuple::empty(py), None).unwrap();
		closure.__call__(PyTuple::empty(py), None).unwrap();
		closure.__call__(PyTuple::empty(py), None).unwrap();
    });
}

It'll be nice to have a more idiomatic interface for that.

src/types/function.rs

LaurentMazare · 2021-10-04T06:36:22Z

Thanks for the feedback and the detailed example showing why Send + 'static are required, this is very helpful.

Re making it more idiomatic having some helper functions for the no-args or no-returns case would be an easy thing to do. Alternatively a dedicated macro similar to #[pyfunction]/wrap_pyfunction! could be helpful there though it would add quite some complexity.
For the return type of the closure, maybe just requiring it to be IntoPy would do the trick, though not sure that there could be an equivalent for the closure arguments.

      pub fn new_closure<F, R>(f: F, py: Python) -> PyResult<&PyCFunction>
      where
          F: Fn(&types::PyTuple, Option<&types::PyDict>) -> PyResult<R> + Send + 'static,
          R: IntoPy<PyObject>,
      {

src/types/function.rs

davidhewitt · 2021-10-04T08:05:09Z

Thanks, I've wanted this for a long time too! There's a very old issue #804. I don't know if we want to take any inspiration from wasm_bindgen as to how we want to implement this?

Alternatively a dedicated macro similar to #[pyfunction]/wrap_pyfunction! could be helpful there though it would add quite some complexity.

I agree that if we wanted to support varied argument types, it would be best to have a proc macro. py_closure! perhaps? We would be able to piggy-back off a lot of the pyfunction codegen, although might force some refactoring.

For the return type of the closure, maybe just requiring it to be IntoPy would do the trick

I think IntoPyCallbackOutput<*mut ffi::PyObject> is probably the best trait to use for the return value. It'll allow anything which implements IntoPy<PyObject>, whether wrapped in a result or not, as well as ().

(This is essentially what #[pyfunction] uses.)

LaurentMazare · 2021-10-04T10:54:14Z

Thanks all for the quick feedback and the reviews. One thing I forgot to mention, mostly for context, is that this wrapping of a PyCapsule in PyCFunction_NewEx is what we use in pyml to expose OCaml closures to the Python runtime.

Thanks, I've wanted this for a long time too! There's a very old issue #804. I don't know if we want to take any inspiration from wasm_bindgen as to how we want to implement this?

Ah sorry I haven't noticed this issue. I didn't know about wasm_bindgen Closure so I only glimpsed at the documentation. My (very basic) understanding is that this brings mostly a way to control the lifetime of the closure on the Rust side and invalidate the wrapping function once it has been dropped. For the use cases that I have at the moment I would have thought that tying the lifetime to the Python function still being reachable is somewhat more natural, but maybe there are other use cases where being able to drop the function from the Rust side is helpful?

I agree that if we wanted to support varied argument types, it would be best to have a proc macro. py_closure! perhaps? We would be able to piggy-back off a lot of the pyfunction codegen, although might force some refactoring.

Makes sense, this seems like a bit of a more involved change so maybe would it be ok to defer it to another PR?

I think IntoPyCallbackOutput<*mut ffi::PyObject> is probably the best trait to use for the return value. It'll allow anything which implements IntoPy<PyObject>, whether wrapped in a result or not, as well as ().

Got it, the PR has been tweaked to use this IntoPyCallbackOutput<*mut ffi::PyObject>, and one of the example has been adapted to benefit from this.

mejrs · 2021-10-04T23:13:59Z

I think it would be better if users can omit the type annotations in the closure arguments - like below. This doesn't compile now:

fn main() {
    Python::with_gil(|py| {
        let counter_fn = |_, _| {
            println!("called");
            Ok(())
        };
	
        let counter_py: Py<PyCFunction> = PyCFunction::new_closure("", counter_fn, py).unwrap().into();
    });
}

error: implementation of `FnOnce` is not general enough
  --> src\main.rs:11:43
   |
11 |         let counter_py: Py<PyCFunction> = PyCFunction::new_closure("", counter_fn, py).unwrap().into();
   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^ implementation of `FnOnce` is not general enough
   |
   = note: closure with signature `fn(&'2 PyTuple, Option<&PyDict>) -> Result<(), PyErr>` must implement `FnOnce<(&'1 PyTuple, Option<&PyDict>)>`, for any lifetime `'1`...
   = note: ...but it actually implements `FnOnce<(&'2 PyTuple, Option<&PyDict>)>`, for some specific lifetime `'2`

error: implementation of `FnOnce` is not general enough
  --> src\main.rs:11:43
   |
11 |         let counter_py: Py<PyCFunction> = PyCFunction::new_closure("", counter_fn, py).unwrap().into();
   |                                           ^^^^^^^^^^^^^^^^^^^^^^^^ implementation of `FnOnce` is not general enough
   |
   = note: closure with signature `fn(&PyTuple, Option<&'2 PyDict>) -> Result<(), PyErr>` must implement `FnOnce<(&PyTuple, Option<&'1 PyDict>)>`, for any lifetime `'1`...
   = note: ...but it actually implements `FnOnce<(&PyTuple, Option<&'2 PyDict>)>`, for some specific lifetime `'2`

Perhaps that can actually work. However, this requires higher rank trait bounds (?) and I couldn't get it to work (safely).

LaurentMazare · 2021-10-05T07:24:47Z

I think it would be better if users can omit the type annotations in the closure arguments - like below.

Agreed that it would be better, though maybe if we have a pyclosure! macro this is less important? The type annotations can be reduced so as just to mention the references but it's not as nice as removing them entirely.

      let counter_fn = move |_: &_, _: Option<&_>| {
          Ok(())
      };

LaurentMazare · 2021-10-11T14:08:51Z

@mejrs @birkenfeld @davidhewitt any more thoughts/feedback on this? Would be pretty keen to move forward with it as I have some current use of pyo3 relying on some class + __call__ wrappers to expose rust functions "dynamically" and this would make it a bit easier.

mejrs

Thanks!

There some things that need doing:

An ui test to check that the use-after-free example above doesn't compile (and keeps not compiling). We use trybuild for that, see https://github.com/PyO3/pyo3/blob/main/tests/test_compile_error.rs
A doc example would be awesome 🙏🏻
You need to fetch upstream changes to make CI play nice (don't worry about the PyPy failures for now)

mejrs · 2021-10-11T16:24:45Z

src/types/function.rs

+    /// Create a new function from a closure.
+    pub fn new_closure<F, R>(f: F, py: Python) -> PyResult<&PyCFunction>
+    where
+        F: Fn(&types::PyTuple, Option<&types::PyDict>) -> PyResult<R> + Send + 'static,


Suggested change

F: Fn(&types::PyTuple, Option<&types::PyDict>) -> PyResult<R> + Send + 'static,

F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,

(same for the other places where this bound is used)

This allows for returning () from the closure, so users don't have to return anything from the closure:

use pyo3::prelude::*; use pyo3::types::{PyCFunction, PyDict, PyTuple}; fn main() { let fun = Python::with_gil(|py| { let fun= |_args: &PyTuple, _kwargs: Option<&PyDict>| { println!("hi"); }; let counter_py: Py<PyCFunction> = PyCFunction::new_closure(fun, py).unwrap().into(); counter_py }); Python::with_gil(|py| { fun.call0(py).unwrap(); }); }

This still allows for returning a result, because of impl<T, E, U> IntoPyCallbackOutput<U> for Result<T, E>

Sounds like a good idea, I've just pushed some changes for this.

LaurentMazare · 2021-10-11T17:55:43Z

@mejrs thanks for the suggestions:

I've added a ui test based on the initial issue you spotted, seem to be reporting the lifetime error properly.
A small doc example has been added too, and fails appropriately if tweaking the assert.
I've made a bit of a mess in the rebase process but hopefully this should be sorted out now.

birkenfeld · 2021-10-13T05:55:34Z

src/types/function.rs

+    {
+        let function_ptr = Box::into_raw(Box::new(f));
+        let capsule_ptr = unsafe {
+            ffi::PyCapsule_New(


This could return NULL on error I assume?

Good catch - it could be worth wrapping this ffi call in Py::from_owned_ptr_or_err, which would also avoid the need for the Py_DECREF at the end.

Nice catch indeed, switched to Py::from_owned_ptr_or_err and removed the decref call (I also checked manually that the closure drop is still called).

davidhewitt

Thanks, this looks overall great. I have a few final comments / questions!

davidhewitt · 2021-10-13T07:18:42Z

src/types/function.rs

+    {
+        let function_ptr = Box::into_raw(Box::new(f));
+        let capsule_ptr = unsafe {
+            ffi::PyCapsule_New(


Good catch - it could be worth wrapping this ffi call in Py::from_owned_ptr_or_err, which would also avoid the need for the Py_DECREF at the end.

davidhewitt · 2021-10-13T07:21:24Z

src/types/function.rs

        method_def: PyMethodDef,
-        py_or_module: PyFunctionArguments,
+        py: Python,
+        mod_ptr: *mut ffi::PyObject,


I see, so in the closure case the mod_ptr is set to the capsule, which makes the capsule object appear as the first argument to run_closure. That's interesting!

davidhewitt · 2021-10-13T07:23:18Z

src/types/function.rs

+    if let Err(err) = result {
+        eprintln!("--- PyO3 intercepted a panic when dropping a closure");


I'm not a huge fan of this but I don't know what else we could do either. There's always the risk that eprintln! can also panic and also lead to UB. So it might need a second std::panic::catch_unwind which aborts if we want to be 100% safe.

Right, I've added another std::panic::catch_unwind layer as you suggest, though this required me to use AssertUnwindSafe which might be wrong if the type returned by the panic had some interior mutability.

davidhewitt · 2021-10-13T07:24:55Z

tests/test_pyfunction.rs

+#[test]
+fn test_closure_counter() {
+    let gil = Python::acquire_gil();
+    let py = gil.python();
+
+    let counter = std::cell::RefCell::new(0);
+    let counter_fn =
+        move |_args: &types::PyTuple, _kwargs: Option<&types::PyDict>| -> PyResult<i32> {
+            let mut counter = counter.borrow_mut();
+            *counter += 1;
+            Ok(*counter)
+        };
+    let counter_py = PyCFunction::new_closure(counter_fn, py).unwrap();
+
+    py_assert!(py, counter_py, "counter_py() == 1");
+    py_assert!(py, counter_py, "counter_py() == 2");
+    py_assert!(py, counter_py, "counter_py() == 3");
+}


This is interesting. I guess because of the GIL we can agree that the RefCell usage is safe? Do you think it would even be safe for us to support FnMut?

I think you're right and holding the GIL should prevent the boxed function to be called in parallel so FnMut should be fine (though I'm not super confident on all this). I've switched to allowing FnMut and also tweaked the example so that it doesn't require a RefCell anymore and instead just uses a Box.

Hmmm, I'm afraid I think I spoke too quickly, sorry. I've tweaked the Box example a bit to one which holds an &mut reference to the Box while calling itself recursively in another thread, which looks to me like it breaks the rule that only one &mut reference to a value can exist at a time. (i.e. it's probably possible to create UB with FnMut.)

let mut counter = Box::new(0); let counter_fn_obj: Arc<RwLock<Option<Py<PyCFunction>>>> = Arc::new(RwLock::new(None)); let counter_fn_obj_clone = counter_fn_obj.clone(); let counter_fn = move |_args: &types::PyTuple, _kwargs: Option<&types::PyDict>| -> PyResult<i32> { let counter_value = &mut *counter; if *counter_value < 5 { _args.py().allow_threads(|| { let counter_fn_obj_clone = counter_fn_obj_clone.clone(); let other_thread = std::thread::spawn(move || Python::with_gil(|py| { counter_fn_obj_clone.read().as_ref().expect("function has been made").call0(py) })); *counter_value += 1; other_thread.join().unwrap() })?; } Ok(*counter_value) }; let counter_py = PyCFunction::new_closure(counter_fn, py).unwrap(); *counter_fn_obj.write() = Some(counter_py.into());

So I think to be safe let's go back to just Fn. Sorry about that.

(I think once we switch back to Fn this looks good to merge in my opinion!)

Oh, thanks for checking this thoroughly, just reverted the FnMut changes.

davidhewitt · 2021-10-17T07:02:22Z

I've just rebased and added a CHANGELOG entry, will proceed to merge this. Thanks again!

LaurentMazare · 2021-10-17T10:15:23Z

Yay, thanks all for the feedback and help getting this out!

mejrs reviewed Oct 4, 2021

View reviewed changes

src/types/function.rs Outdated Show resolved Hide resolved

src/types/function.rs Outdated Show resolved Hide resolved

birkenfeld reviewed Oct 4, 2021

View reviewed changes

src/types/function.rs Outdated Show resolved Hide resolved

src/types/function.rs Outdated Show resolved Hide resolved

src/types/function.rs Outdated Show resolved Hide resolved

src/types/function.rs Outdated Show resolved Hide resolved

src/types/function.rs Show resolved Hide resolved

mejrs requested changes Oct 11, 2021

View reviewed changes

LaurentMazare force-pushed the closures branch 2 times, most recently from 1723218 to 57f9ba2 Compare October 11, 2021 17:00

mejrs approved these changes Oct 11, 2021

View reviewed changes

birkenfeld reviewed Oct 13, 2021

View reviewed changes

davidhewitt reviewed Oct 13, 2021

View reviewed changes

pycfunction: add PyCFunction::new_closure rust defined closures

2042906

davidhewitt force-pushed the closures branch from 68c1932 to 2042906 Compare October 17, 2021 07:00

davidhewitt enabled auto-merge October 17, 2021 07:02

davidhewitt merged commit fbb5e3c into PyO3:main Oct 17, 2021

davidhewitt mentioned this pull request Aug 21, 2022

Wasm-bindgen-like Closure wrapper around Rust functions? #804

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support for wrapping rust closures as python functions #1901

Support for wrapping rust closures as python functions #1901

LaurentMazare commented Oct 3, 2021

mejrs left a comment

LaurentMazare commented Oct 4, 2021

davidhewitt commented Oct 4, 2021

LaurentMazare commented Oct 4, 2021

mejrs commented Oct 4, 2021 •

edited

LaurentMazare commented Oct 5, 2021

LaurentMazare commented Oct 11, 2021

mejrs left a comment

mejrs Oct 11, 2021

LaurentMazare Oct 11, 2021

LaurentMazare Oct 11, 2021

LaurentMazare commented Oct 11, 2021

birkenfeld Oct 13, 2021

davidhewitt Oct 13, 2021

LaurentMazare Oct 13, 2021

davidhewitt left a comment

davidhewitt Oct 13, 2021

davidhewitt Oct 13, 2021

davidhewitt Oct 13, 2021

LaurentMazare Oct 13, 2021

davidhewitt Oct 13, 2021

LaurentMazare Oct 13, 2021

davidhewitt Oct 13, 2021

davidhewitt Oct 13, 2021

LaurentMazare Oct 14, 2021

davidhewitt commented Oct 17, 2021 •

edited

LaurentMazare commented Oct 17, 2021

	F: Fn(&types::PyTuple, Option<&types::PyDict>) -> PyResult<R> + Send + 'static,
	F: Fn(&types::PyTuple, Option<&types::PyDict>) -> R + Send + 'static,

		if let Err(err) = result {
		eprintln!("--- PyO3 intercepted a panic when dropping a closure");

Support for wrapping rust closures as python functions #1901

Support for wrapping rust closures as python functions #1901

Conversation

LaurentMazare commented Oct 3, 2021

mejrs left a comment

Choose a reason for hiding this comment

LaurentMazare commented Oct 4, 2021

davidhewitt commented Oct 4, 2021

LaurentMazare commented Oct 4, 2021

mejrs commented Oct 4, 2021 • edited

LaurentMazare commented Oct 5, 2021

LaurentMazare commented Oct 11, 2021

mejrs left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

LaurentMazare commented Oct 11, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davidhewitt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davidhewitt commented Oct 17, 2021 • edited

LaurentMazare commented Oct 17, 2021

mejrs commented Oct 4, 2021 •

edited

davidhewitt commented Oct 17, 2021 •

edited