Add BoundPyFunction helper struct

This standardises the return value of `_wrap_pyfunction` and `_wrap_pyfunction_bound`.
PyO3 · Feb 25, 2024 · 2f8f123 · 2f8f123
1 parent 500b114
commit 2f8f123
Show file tree

Hide file tree

Showing 5 changed files with 91 additions and 28 deletions.
diff --git a/pyo3-macros-backend/src/module.rs b/pyo3-macros-backend/src/module.rs
@@ -272,7 +272,7 @@ fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn
                 let name = &func.sig.ident;
                 let statements: Vec<syn::Stmt> = syn::parse_quote! {
                     #wrapped_function
-                    #module_name.add_function(#krate::impl_::pyfunction::_wrap_pyfunction(&#name::DEF, #module_name)?)?;
+                    #module_name.add_function(#krate::impl_::pyfunction::_wrap_pyfunction(&#name::DEF, #module_name)?.into())?;
                 };
                 stmts.extend(statements);
             }

diff --git a/src/impl_/pyfunction.rs b/src/impl_/pyfunction.rs
@@ -1,21 +1,44 @@
 use crate::{
     derive_utils::{PyFunctionArguments, PyFunctionArgumentsBound},
     types::PyCFunction,
-    Bound, PyResult,
+    Bound, Py, PyResult,
 };
 
 pub use crate::impl_::pymethods::PyMethodDef;
 
 pub fn _wrap_pyfunction<'a>(
     method_def: &PyMethodDef,
     py_or_module: impl Into<PyFunctionArguments<'a>>,
-) -> PyResult<&'a PyCFunction> {
-    PyCFunction::internal_new(method_def, py_or_module.into()).map(|x| x.into_gil_ref())
+) -> PyResult<BoundPyFunction<'a>> {
+    PyCFunction::internal_new(method_def, py_or_module.into()).map(BoundPyFunction)
 }
 
 pub fn _wrap_pyfunction_bound<'a, 'py: 'a>(
     method_def: &PyMethodDef,
     py_or_module: impl Into<PyFunctionArgumentsBound<'a, 'py>>,
-) -> PyResult<Bound<'py, PyCFunction>> {
-    PyCFunction::internal_new_bound(method_def, py_or_module.into())
+) -> PyResult<BoundPyFunction<'py>> {
+    PyCFunction::internal_new_bound(method_def, py_or_module.into()).map(BoundPyFunction)
+}
+
+pub struct BoundPyFunction<'py>(pub Bound<'py, PyCFunction>);
+
+impl<'a> From<BoundPyFunction<'a>> for &'a PyCFunction {
+    #[inline]
+    fn from(bound: BoundPyFunction<'a>) -> Self {
+        bound.0.into_gil_ref()
+    }
+}
+
+impl<'py> From<BoundPyFunction<'py>> for Bound<'py, PyCFunction> {
+    #[inline]
+    fn from(bound: BoundPyFunction<'py>) -> Self {
+        bound.0
+    }
+}
+
+impl<'py> From<BoundPyFunction<'py>> for Py<PyCFunction> {
+    #[inline]
+    fn from(bound: BoundPyFunction<'py>) -> Self {
+        bound.0.unbind()
+    }
 }
diff --git a/src/macros.rs b/src/macros.rs
@@ -126,13 +126,25 @@ macro_rules! py_run_impl {
 macro_rules! wrap_pyfunction {
     ($function:path) => {
         &|py_or_module| {
+            use ::std::convert::Into;
+            use $crate::types::PyCFunction;
+            use $crate::PyResult;
             use $function as wrapped_pyfunction;
-            $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, py_or_module)
+            let wrapped: PyResult<&PyCFunction> =
+                $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, py_or_module)
+                    .map(|f| f.into());
+            wrapped
         }
     };
     ($function:path, $py_or_module:expr) => {{
+        use ::std::convert::Into;
+        use $crate::types::PyCFunction;
+        use $crate::PyResult;
         use $function as wrapped_pyfunction;
-        $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, $py_or_module)
+        let wrapped: PyResult<&PyCFunction> =
+            $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, $py_or_module)
+                .map(|f| f.into());
+        wrapped
     }};
 }
 
@@ -144,16 +156,31 @@ macro_rules! wrap_pyfunction {
 macro_rules! wrap_pyfunction_bound {
     ($function:path) => {
         &|py_or_module| {
+            use ::std::convert::Into;
+            use $crate::types::PyCFunction;
+            use $crate::PyResult;
             use $function as wrapped_pyfunction;
-            $crate::impl_::pyfunction::_wrap_pyfunction_bound(
-                &wrapped_pyfunction::DEF,
-                py_or_module,
-            )
+            let wrapped: PyResult<Bound<'_, PyCFunction>> =
+                $crate::impl_::pyfunction::_wrap_pyfunction_bound(
+                    &wrapped_pyfunction::DEF,
+                    py_or_module,
+                )
+                .map(|f| f.into());
+            wrapped
         }
     };
     ($function:path, $py_or_module:expr) => {{
+        use ::std::convert::Into;
+        use $crate::types::PyCFunction;
+        use $crate::PyResult;
         use $function as wrapped_pyfunction;
-        $crate::impl_::pyfunction::_wrap_pyfunction_bound(&wrapped_pyfunction::DEF, $py_or_module)
+        let wrapped: PyResult<Bound<'_, PyCFunction>> =
+            $crate::impl_::pyfunction::_wrap_pyfunction_bound(
+                &wrapped_pyfunction::DEF,
+                $py_or_module,
+            )
+            .map(|f| f.into());
+        wrapped
     }};
 }
 

diff --git a/src/types/function.rs b/src/types/function.rs
@@ -34,17 +34,25 @@ impl PyCFunction {
         doc: &'static str,
         py_or_module: PyFunctionArguments<'a>,
     ) -> PyResult<&'a Self> {
-        Self::new_with_keywords_bound(fun, name, doc, py_or_module).map(Bound::into_gil_ref)
+        Self::internal_new(
+            &PyMethodDef::cfunction_with_keywords(
+                name,
+                pymethods::PyCFunctionWithKeywords(fun),
+                doc,
+            ),
+            py_or_module,
+        )
+        .map(Bound::into_gil_ref)
     }
 
     /// Create a new built-in function with keywords (*args and/or **kwargs).
-    pub fn new_with_keywords_bound<'a>(
+    pub fn new_with_keywords_bound<'py>(
         fun: ffi::PyCFunctionWithKeywords,
         name: &'static str,
         doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
-        Self::internal_new(
+        py_or_module: PyFunctionArgumentsBound<'_, 'py>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
             &PyMethodDef::cfunction_with_keywords(
                 name,
                 pymethods::PyCFunctionWithKeywords(fun),
@@ -68,17 +76,21 @@ impl PyCFunction {
         doc: &'static str,
         py_or_module: PyFunctionArguments<'a>,
     ) -> PyResult<&'a Self> {
-        Self::new_bound(fun, name, doc, py_or_module).map(Bound::into_gil_ref)
+        Self::internal_new(
+            &PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
+            py_or_module,
+        )
+        .map(Bound::into_gil_ref)
     }
 
     /// Create a new built-in function which takes no arguments.
-    pub fn new_bound<'a>(
+    pub fn new_bound<'py>(
         fun: ffi::PyCFunction,
         name: &'static str,
         doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
-        Self::internal_new(
+        py_or_module: PyFunctionArgumentsBound<'_, 'py>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
             &PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
             py_or_module,
         )

diff --git a/tests/test_module.rs b/tests/test_module.rs
@@ -295,18 +295,19 @@ fn test_module_nesting() {
 // Test that argument parsing specification works for pyfunctions
 
 #[pyfunction(signature = (a=5, *args))]
-fn ext_vararg_fn(py: Python<'_>, a: i32, args: &PyTuple) -> PyObject {
-    [a.to_object(py), args.into()].to_object(py)
+fn ext_vararg_fn(py: Python<'_>, a: i32, args: &Bound<'_, PyTuple>) -> PyObject {
+    [a.to_object(py), args.into_py(py)].to_object(py)
 }
 
 #[pymodule]
-fn vararg_module(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+fn vararg_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     #[pyfn(m, signature = (a=5, *args))]
-    fn int_vararg_fn(py: Python<'_>, a: i32, args: &PyTuple) -> PyObject {
+    fn int_vararg_fn(py: Python<'_>, a: i32, args: &Bound<'_, PyTuple>) -> PyObject {
         ext_vararg_fn(py, a, args)
     }
 
-    m.add_function(wrap_pyfunction!(ext_vararg_fn, m)?).unwrap();
+    m.add_function(&wrap_pyfunction_bound!(ext_vararg_fn, m)?)
+        .unwrap();
     Ok(())
 }