Pymodule bound (#3897)

* Support Bound in pymodule and pyfunction macros

Co-authored-by: David Hewitt <mail@davidhewitt.dev>

* Remove spurious $ character

Co-authored-by: Matthew Neeley <mneeley@gmail.com>

* Rework PyCFunction::new_bound signatures

This allows us to remove the awkward `PyFunctionArgumentsBound` enum.

* Use BoundRef instead of BoundModule

* support argument deduction for `wrap_pyfunction_bound!`

* support `wrap_pyfunction!` with `Bound` input/output

* Fix docs link to `wrap_pyfunction_bound!`

* Revert back to wrap_pyfunction!

---------

Co-authored-by: David Hewitt <mail@davidhewitt.dev>
Co-authored-by: Matthew Neeley <mneeley@gmail.com>

guide/src/module.md

@@ -79,7 +79,7 @@ fn parent_module(py: Python<'_>, m: &PyModule) -> PyResult<()> {
 
 fn register_child_module(py: Python<'_>, parent_module: &PyModule) -> PyResult<()> {
     let child_module = PyModule::new_bound(py, "child_module")?;
-    child_module.add_function(&wrap_pyfunction!(func, child_module.as_gil_ref())?.as_borrowed())?;
+    child_module.add_function(wrap_pyfunction!(func, &child_module)?)?;
     parent_module.add_submodule(child_module.as_gil_ref())?;
     Ok(())
 }

guide/src/python_from_rust.md

@@ -311,7 +311,7 @@ fn main() -> PyResult<()> {
     Python::with_gil(|py| {
         // Create new module
         let foo_module = PyModule::new_bound(py, "foo")?;
-        foo_module.add_function(&wrap_pyfunction!(add_one, foo_module.as_gil_ref())?.as_borrowed())?;
+        foo_module.add_function(wrap_pyfunction!(add_one, &foo_module)?)?;
 
         // Import and get sys.modules
         let sys = PyModule::import_bound(py, "sys")?;

pyo3-macros-backend/src/module.rs

@@ -215,15 +215,16 @@ pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream>
         #[allow(unknown_lints, non_local_definitions)]
         const _: () = {
             use #krate::impl_::pymodule as impl_;
+            use #krate::impl_::pymethods::BoundRef;
 
             fn __pyo3_pymodule(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
-                #ident(module.py(), module.as_gil_ref())
+                #ident(module.py(), ::std::convert::Into::into(BoundRef(module)))
             }
 
             impl #ident::MakeDef {
                 const fn make_def() -> impl_::ModuleDef {
+                    const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
                     unsafe {
-                        const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
                         impl_::ModuleDef::new(
                             #ident::__PYO3_NAME,
                             #doc,
@@ -263,9 +264,13 @@ fn module_initialization(options: PyModuleOptions, ident: &syn::Ident) -> TokenS
 
 /// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
 fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn) -> Result<()> {
-    let mut stmts: Vec<syn::Stmt> = Vec::new();
     let krate = get_pyo3_crate(&options.krate);
 
+    let mut stmts: Vec<syn::Stmt> = vec![syn::parse_quote!(
+        #[allow(unknown_lints, unused_imports, redundant_imports)]
+        use #krate::{PyNativeType, types::PyModuleMethods};
+    )];
+
     for mut stmt in func.block.stmts.drain(..) {
         if let syn::Stmt::Item(Item::Fn(func)) = &mut stmt {
             if let Some(pyfn_args) = get_pyfn_attr(&mut func.attrs)? {
@@ -274,7 +279,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.as_borrowed().add_function(#krate::wrap_pyfunction!(#name, #module_name.as_borrowed())?)?;
                 };
                 stmts.extend(statements);
             }

pyo3-macros-backend/src/pyfunction.rs

@@ -268,12 +268,12 @@ pub fn impl_wrap_pyfunction(
         #[doc(hidden)]
         #vis mod #name {
             pub(crate) struct MakeDef;
-            pub const DEF: #krate::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
+            pub const DEF: #krate::impl_::pymethods::PyMethodDef = MakeDef::DEF;
 
             pub fn add_to_module(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
                 use #krate::prelude::PyModuleMethods;
                 use ::std::convert::Into;
-                module.add_function(&#krate::types::PyCFunction::internal_new(&DEF, module.as_gil_ref().into())?)
+                module.add_function(#krate::types::PyCFunction::internal_new(&DEF, module.as_gil_ref().into())?)
             }
         }
 
@@ -287,7 +287,7 @@ pub fn impl_wrap_pyfunction(
             use #krate as _pyo3;
 
             impl #name::MakeDef {
-                const DEF: #krate::impl_::pyfunction::PyMethodDef = #methoddef;
+                const DEF: #krate::impl_::pymethods::PyMethodDef = #methoddef;
             }
 
             #[allow(non_snake_case)]

src/impl_/pyfunction.rs

@@ -1,10 +1,68 @@
-use crate::{derive_utils::PyFunctionArguments, types::PyCFunction, PyResult};
+use crate::{
+    types::{PyCFunction, PyModule},
+    Borrowed, Bound, PyResult, Python,
+};
 
 pub use crate::impl_::pymethods::PyMethodDef;
 
-pub fn _wrap_pyfunction<'a>(
+/// Trait to enable the use of `wrap_pyfunction` with both `Python` and `PyModule`,
-    method_def: &PyMethodDef,
+/// and also to infer the return type of either `&'py PyCFunction` or `Bound<'py, PyCFunction>`.
-    py_or_module: impl Into<PyFunctionArguments<'a>>,
+pub trait WrapPyFunctionArg<'py, T> {
-) -> PyResult<&'a PyCFunction> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<T>;
-    PyCFunction::internal_new(method_def, py_or_module.into()).map(|x| x.into_gil_ref())
+}
+
+impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for Bound<'py, PyModule> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        PyCFunction::internal_new_bound(self.py(), method_def, Some(&self))
+    }
+}
+
+impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for &'_ Bound<'py, PyModule> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        PyCFunction::internal_new_bound(self.py(), method_def, Some(self))
+    }
+}
+
+impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for Borrowed<'_, 'py, PyModule> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        PyCFunction::internal_new_bound(self.py(), method_def, Some(&self))
+    }
+}
+
+impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for &'_ Borrowed<'_, 'py, PyModule> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        PyCFunction::internal_new_bound(self.py(), method_def, Some(self))
+    }
+}
+
+// For Python<'py>, only the GIL Ref form exists to avoid causing type inference to kick in.
+// The `wrap_pyfunction_bound!` macro is needed for the Bound form.
+impl<'py> WrapPyFunctionArg<'py, &'py PyCFunction> for Python<'py> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<&'py PyCFunction> {
+        PyCFunction::internal_new(method_def, self.into()).map(Bound::into_gil_ref)
+    }
+}
+
+impl<'py> WrapPyFunctionArg<'py, &'py PyCFunction> for &'py PyModule {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<&'py PyCFunction> {
+        PyCFunction::internal_new(method_def, self.into()).map(Bound::into_gil_ref)
+    }
+}
+
+/// Helper for `wrap_pyfunction_bound!` to guarantee return type of `Bound<'py, PyCFunction>`.
+pub struct OnlyBound<T>(pub T);
+
+impl<'py, T> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for OnlyBound<T>
+where
+    T: WrapPyFunctionArg<'py, Bound<'py, PyCFunction>>,
+{
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        WrapPyFunctionArg::wrap_pyfunction(self.0, method_def)
+    }
+}
+
+impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for OnlyBound<Python<'py>> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
+        PyCFunction::internal_new_bound(self.0, method_def, None)
+    }
 }

src/macros.rs

@@ -121,18 +121,57 @@ macro_rules! py_run_impl {
 /// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
 ///
 /// This can be used with [`PyModule::add_function`](crate::types::PyModule::add_function) to add free
-/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more information.
+/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more
+/// information.
+///
+/// During the migration from the GIL Ref API to the Bound API, the return type of this macro will
+/// be either the `&'py PyModule` GIL Ref or `Bound<'py, PyModule>` according to the second
+/// argument.
+///
+/// For backwards compatibility, if the second argument is `Python<'py>` then the return type will
+/// be `&'py PyModule` GIL Ref. To get `Bound<'py, PyModule>`, use the [`crate::wrap_pyfunction_bound!`]
+/// macro instead.
 #[macro_export]
 macro_rules! wrap_pyfunction {
     ($function:path) => {
         &|py_or_module| {
             use $function as wrapped_pyfunction;
-            $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, py_or_module)
+            $crate::impl_::pyfunction::WrapPyFunctionArg::wrap_pyfunction(
+                py_or_module,
+                &wrapped_pyfunction::DEF,
+            )
         }
     };
     ($function:path, $py_or_module:expr) => {{
         use $function as wrapped_pyfunction;
-        $crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, $py_or_module)
+        $crate::impl_::pyfunction::WrapPyFunctionArg::wrap_pyfunction(
+            $py_or_module,
+            &wrapped_pyfunction::DEF,
+        )
+    }};
+}
+
+/// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
+///
+/// This can be used with [`PyModule::add_function`](crate::types::PyModule::add_function) to add free
+/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more information.
+#[macro_export]
+macro_rules! wrap_pyfunction_bound {
+    ($function:path) => {
+        &|py_or_module| {
+            use $function as wrapped_pyfunction;
+            $crate::impl_::pyfunction::WrapPyFunctionArg::wrap_pyfunction(
+                $crate::impl_::pyfunction::OnlyBound(py_or_module),
+                &wrapped_pyfunction::DEF,
+            )
+        }
+    };
+    ($function:path, $py_or_module:expr) => {{
+        use $function as wrapped_pyfunction;
+        $crate::impl_::pyfunction::WrapPyFunctionArg::wrap_pyfunction(
+            $crate::impl_::pyfunction::OnlyBound($py_or_module),
+            &wrapped_pyfunction::DEF,
+        )
     }};
 }
 

src/prelude.rs

@@ -23,7 +23,7 @@ pub use crate::PyNativeType;
 pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, FromPyObject};
 
 #[cfg(feature = "macros")]
-pub use crate::wrap_pyfunction;
+pub use crate::{wrap_pyfunction, wrap_pyfunction_bound};
 
 pub use crate::types::any::PyAnyMethods;
 pub use crate::types::boolobject::PyBoolMethods;

src/types/function.rs

@@ -3,10 +3,11 @@ use crate::ffi_ptr_ext::FfiPtrExt;
 use crate::methods::PyMethodDefDestructor;
 use crate::py_result_ext::PyResultExt;
 use crate::types::capsule::PyCapsuleMethods;
+use crate::types::module::PyModuleMethods;
 use crate::{
     ffi,
     impl_::pymethods::{self, PyMethodDef},
-    types::{PyCapsule, PyDict, PyString, PyTuple},
+    types::{PyCapsule, PyDict, PyModule, PyString, PyTuple},
 };
 use crate::{Bound, IntoPy, Py, PyAny, PyResult, Python};
 use std::cell::UnsafeCell;
@@ -33,16 +34,6 @@ 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)
-    }
-
-    /// Create a new built-in function with keywords (*args and/or **kwargs).
-    pub fn new_with_keywords_bound<'a>(
-        fun: ffi::PyCFunctionWithKeywords,
-        name: &'static str,
-        doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
         Self::internal_new(
             &PyMethodDef::cfunction_with_keywords(
                 name,
@@ -51,6 +42,26 @@ impl PyCFunction {
             ),
             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<'py>(
+        py: Python<'py>,
+        fun: ffi::PyCFunctionWithKeywords,
+        name: &'static str,
+        doc: &'static str,
+        module: Option<&Bound<'py, PyModule>>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
+            py,
+            &PyMethodDef::cfunction_with_keywords(
+                name,
+                pymethods::PyCFunctionWithKeywords(fun),
+                doc,
+            ),
+            module,
+        )
     }
 
     /// Deprecated form of [`PyCFunction::new`]
@@ -67,20 +78,26 @@ 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)
-    }
-
-    /// Create a new built-in function which takes no arguments.
-    pub fn new_bound<'a>(
-        fun: ffi::PyCFunction,
-        name: &'static str,
-        doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
         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<'py>(
+        py: Python<'py>,
+        fun: ffi::PyCFunction,
+        name: &'static str,
+        doc: &'static str,
+        module: Option<&Bound<'py, PyModule>>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
+            py,
+            &PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
+            module,
+        )
     }
 
     /// Deprecated form of [`PyCFunction::new_closure`]
@@ -189,6 +206,35 @@ impl PyCFunction {
                 .downcast_into_unchecked()
         }
     }
+
+    #[doc(hidden)]
+    pub(crate) fn internal_new_bound<'py>(
+        py: Python<'py>,
+        method_def: &PyMethodDef,
+        module: Option<&Bound<'py, PyModule>>,
+    ) -> PyResult<Bound<'py, Self>> {
+        let (mod_ptr, module_name): (_, Option<Py<PyString>>) = if let Some(m) = module {
+            let mod_ptr = m.as_ptr();
+            (mod_ptr, Some(m.name()?.into_py(py)))
+        } else {
+            (std::ptr::null_mut(), None)
+        };
+        let (def, destructor) = method_def.as_method_def()?;
+
+        // FIXME: stop leaking the def and destructor
+        let def = Box::into_raw(Box::new(def));
+        std::mem::forget(destructor);
+
+        let module_name_ptr = module_name
+            .as_ref()
+            .map_or(std::ptr::null_mut(), Py::as_ptr);
+
+        unsafe {
+            ffi::PyCFunction_NewEx(def, mod_ptr, module_name_ptr)
+                .assume_owned_or_err(py)
+                .downcast_into_unchecked()
+        }
+    }
 }
 
 fn closure_capsule_name() -> &'static CStr {

src/types/module.rs

@@ -399,7 +399,8 @@ impl PyModule {
     /// [1]: crate::prelude::pyfunction
     /// [2]: crate::wrap_pyfunction
     pub fn add_function<'a>(&'a self, fun: &'a PyCFunction) -> PyResult<()> {
-        self.as_borrowed().add_function(&fun.as_borrowed())
+        let name = fun.getattr(__name__(self.py()))?.extract()?;
+        self.add(name, fun)
     }
 }
 
@@ -590,7 +591,7 @@ pub trait PyModuleMethods<'py> {
     ///
     /// [1]: crate::prelude::pyfunction
     /// [2]: crate::wrap_pyfunction
-    fn add_function(&self, fun: &Bound<'_, PyCFunction>) -> PyResult<()>;
+    fn add_function(&self, fun: Bound<'_, PyCFunction>) -> PyResult<()>;
 }
 
 impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
@@ -700,7 +701,7 @@ impl<'py> PyModuleMethods<'py> for Bound<'py, PyModule> {
         self.add(name, module)
     }
 
-    fn add_function(&self, fun: &Bound<'_, PyCFunction>) -> PyResult<()> {
+    fn add_function(&self, fun: Bound<'_, PyCFunction>) -> PyResult<()> {
         let name = fun.getattr(__name__(self.py()))?;
         self.add(name.downcast_into::<PyString>()?, fun)
     }

tests/test_module.rs

@@ -240,7 +240,7 @@ fn subfunction() -> String {
 }
 
 fn submodule(module: &Bound<'_, PyModule>) -> PyResult<()> {
-    module.add_function(&wrap_pyfunction!(subfunction, module.as_gil_ref())?.as_borrowed())?;
+    module.add_function(wrap_pyfunction!(subfunction, module)?)?;
     Ok(())
 }
 
@@ -295,14 +295,14 @@ 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 {
+fn ext_vararg_fn(py: Python<'_>, a: i32, args: &Bound<'_, PyTuple>) -> PyObject {
-    [a.to_object(py), args.into()].to_object(py)
+    [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)
     }
 
@@ -410,7 +410,7 @@ fn pyfunction_with_pass_module_in_attribute(module: &PyModule) -> PyResult<&str>
 }
 
 #[pymodule]
-fn module_with_functions_with_module(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
+fn module_with_functions_with_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(pyfunction_with_module, m)?)?;
     m.add_function(wrap_pyfunction!(pyfunction_with_module_gil_ref, m)?)?;
     m.add_function(wrap_pyfunction!(pyfunction_with_module_owned, m)?)?;

tests/test_pyfunction.rs

@@ -324,10 +324,11 @@ fn test_pycfunction_new() {
         }
 
         let py_fn = PyCFunction::new_bound(
+            py,
             c_fn,
             "py_fn",
             "py_fn for test (this is the docstring)",
-            py.into(),
+            None,
         )
         .unwrap();
 
@@ -381,10 +382,11 @@ fn test_pycfunction_new_with_keywords() {
         }
 
         let py_fn = PyCFunction::new_with_keywords_bound(
+            py,
             c_fn,
             "py_fn",
             "py_fn for test (this is the docstring)",
-            py.into(),
+            None,
         )
         .unwrap();
 

tests/test_wrap_pyfunction_deduction.rs

@@ -13,3 +13,12 @@ pub fn add_wrapped(wrapper: &impl Fn(Python<'_>) -> PyResult<&PyCFunction>) {
 fn wrap_pyfunction_deduction() {
     add_wrapped(wrap_pyfunction!(f));
 }
+
+pub fn add_wrapped_bound(wrapper: &impl Fn(Python<'_>) -> PyResult<Bound<'_, PyCFunction>>) {
+    let _ = wrapper;
+}
+
+#[test]
+fn wrap_pyfunction_deduction_bound() {
+    add_wrapped_bound(wrap_pyfunction_bound!(f));
+}