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>
This commit is contained in:
Lily Foote 2024-02-27 19:19:52 +00:00 committed by GitHub
parent 93704047a5
commit a3ad28b70c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 211 additions and 51 deletions

View File

@ -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(())
}

View File

@ -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")?;

View File

@ -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);
}

View File

@ -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)]

View File

@ -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)
}
}

View File

@ -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,
)
}};
}

View File

@ -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;

View File

@ -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 {

View File

@ -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)
}

View File

@ -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)?)?;

View File

@ -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();

View File

@ -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));
}