Pymodule bound (#3897)

* Support Bound in pymodule and pyfunction macros Co-authored-by: David Hewitt <[email protected]> * Remove spurious $ character Co-authored-by: Matthew Neeley <[email protected]> * 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 <[email protected]> Co-authored-by: Matthew Neeley <[email protected]>
PyO3 · Feb 27, 2024 · a3ad28b · a3ad28b
1 parent 9370404
commit a3ad28b
Show file tree

Hide file tree

Showing 12 changed files with 203 additions and 43 deletions.
diff --git a/guide/src/module.md b/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(())
 }

diff --git a/guide/src/python_from_rust.md b/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")?;

diff --git a/pyo3-macros-backend/src/module.rs b/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);
             }

diff --git a/pyo3-macros-backend/src/pyfunction.rs b/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)]

diff --git a/src/impl_/pyfunction.rs b/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>(
-    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())
+/// Trait to enable the use of `wrap_pyfunction` with both `Python` and `PyModule`,
+/// and also to infer the return type of either `&'py PyCFunction` or `Bound<'py, PyCFunction>`.
+pub trait WrapPyFunctionArg<'py, T> {
+    fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<T>;
+}
+
+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)
+    }
 }
diff --git a/src/macros.rs b/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::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::_wrap_pyfunction(&wrapped_pyfunction::DEF, $py_or_module)
+        $crate::impl_::pyfunction::WrapPyFunctionArg::wrap_pyfunction(
+            $crate::impl_::pyfunction::OnlyBound($py_or_module),
+            &wrapped_pyfunction::DEF,
+        )
     }};
 }
 

diff --git a/src/prelude.rs b/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;

diff --git a/src/types/function.rs b/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,23 +34,33 @@ 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>(
+        py: Python<'py>,
         fun: ffi::PyCFunctionWithKeywords,
         name: &'static str,
         doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
-        Self::internal_new(
+        module: Option<&Bound<'py, PyModule>>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
+            py,
             &PyMethodDef::cfunction_with_keywords(
                 name,
                 pymethods::PyCFunctionWithKeywords(fun),
                 doc,
             ),
-            py_or_module,
+            module,
         )
     }
 
@@ -67,19 +78,25 @@ 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>(
+        py: Python<'py>,
         fun: ffi::PyCFunction,
         name: &'static str,
         doc: &'static str,
-        py_or_module: PyFunctionArguments<'a>,
-    ) -> PyResult<Bound<'a, Self>> {
-        Self::internal_new(
+        module: Option<&Bound<'py, PyModule>>,
+    ) -> PyResult<Bound<'py, Self>> {
+        Self::internal_new_bound(
+            py,
             &PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
-            py_or_module,
+            module,
         )
     }
 
@@ -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 {

diff --git a/src/types/module.rs b/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)
     }