From a3ad28b70ca847459b3f5d8b8644beb826268cc2 Mon Sep 17 00:00:00 2001 From: Lily Foote Date: Tue, 27 Feb 2024 19:19:52 +0000 Subject: [PATCH] Pymodule bound (#3897) * Support Bound in pymodule and pyfunction macros Co-authored-by: David Hewitt * Remove spurious $ character Co-authored-by: Matthew Neeley * 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 Co-authored-by: Matthew Neeley --- guide/src/module.md | 2 +- guide/src/python_from_rust.md | 2 +- pyo3-macros-backend/src/module.rs | 13 ++-- pyo3-macros-backend/src/pyfunction.rs | 6 +- src/impl_/pyfunction.rs | 70 ++++++++++++++++++-- src/macros.rs | 45 ++++++++++++- src/prelude.rs | 2 +- src/types/function.rs | 88 +++++++++++++++++++------ src/types/module.rs | 7 +- tests/test_module.rs | 12 ++-- tests/test_pyfunction.rs | 6 +- tests/test_wrap_pyfunction_deduction.rs | 9 +++ 12 files changed, 211 insertions(+), 51 deletions(-) diff --git a/guide/src/module.md b/guide/src/module.md index 789c3d91..53c390be 100644 --- 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 index 973f997c..b51bbe59 100644 --- 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 index 148cea6f..40cf34f2 100644 --- 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 #[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 = Vec::new(); let krate = get_pyo3_crate(&options.krate); + let mut stmts: Vec = 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::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 index 38362d08..265a4824 100644 --- 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 index 464beeb8..531e0c93 100644 --- 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>, -) -> 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; +} + +impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for Bound<'py, PyModule> { + fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult> { + 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> { + 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> { + 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> { + 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(pub T); + +impl<'py, T> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for OnlyBound +where + T: WrapPyFunctionArg<'py, Bound<'py, PyCFunction>>, +{ + fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult> { + WrapPyFunctionArg::wrap_pyfunction(self.0, method_def) + } +} + +impl<'py> WrapPyFunctionArg<'py, Bound<'py, PyCFunction>> for OnlyBound> { + fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult> { + PyCFunction::internal_new_bound(self.0, method_def, None) + } } diff --git a/src/macros.rs b/src/macros.rs index 9b0d2816..33f378e7 100644 --- 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::_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, + ) }}; } diff --git a/src/prelude.rs b/src/prelude.rs index 1de7c3ac..dcf4fe71 100644 --- 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 index 7cbb05e2..75173d2a 100644 --- 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,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> { 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> { + 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> { 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> { + 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> { + let (mod_ptr, module_name): (_, Option>) = 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 index 245d38d9..75fc6625 100644 --- 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::()?, fun) } diff --git a/tests/test_module.rs b/tests/test_module.rs index 9d14f243..5763043e 100644 --- a/tests/test_module.rs +++ b/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 { - [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) } @@ -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)?)?; diff --git a/tests/test_pyfunction.rs b/tests/test_pyfunction.rs index 14748418..838f7353 100644 --- a/tests/test_pyfunction.rs +++ b/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(); diff --git a/tests/test_wrap_pyfunction_deduction.rs b/tests/test_wrap_pyfunction_deduction.rs index 8723ad43..52c9adcb 100644 --- a/tests/test_wrap_pyfunction_deduction.rs +++ b/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>) { + let _ = wrapper; +} + +#[test] +fn wrap_pyfunction_deduction_bound() { + add_wrapped_bound(wrap_pyfunction_bound!(f)); +}