Skip to content

Commit 031f033

Browse files
davidhewittLilyAcorn
davidhewitt
authored and
LilyAcorn
committed
Support Bound in pymodule and pyfunction macros
1 parent e0e3981 commit 031f033

File tree

13 files changed

+186
-46
lines changed

13 files changed

+186
-46
lines changed

‎guide/src/module.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ fn parent_module(py: Python<'_>, m: &PyModule) -> PyResult<()> {
7979

8080
fn register_child_module(py: Python<'_>, parent_module: &PyModule) -> PyResult<()> {
8181
let child_module = PyModule::new_bound(py, "child_module")?;
82-
child_module.add_function(&wrap_pyfunction!(func, child_module.as_gil_ref())?.as_borrowed())?;
82+
child_module.add_function(wrap_pyfunction_bound!(func, &child_module)?)?;
8383
parent_module.add_submodule(child_module.as_gil_ref())?;
8484
Ok(())
8585
}

‎guide/src/python_from_rust.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ fn main() -> PyResult<()> {
311311
Python::with_gil(|py| {
312312
// Create new module
313313
let foo_module = PyModule::new_bound(py, "foo")?;
314-
foo_module.add_function(&wrap_pyfunction!(add_one, foo_module.as_gil_ref())?.as_borrowed())?;
314+
foo_module.add_function(wrap_pyfunction_bound!(add_one, &foo_module)?)?;
315315

316316
// Import and get sys.modules
317317
let sys = PyModule::import_bound(py, "sys")?;

‎pyo3-macros-backend/src/module.rs

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ pub fn pymodule_module_impl(mut module: syn::ItemMod) -> Result<TokenStream> {
194194
/// module
195195
pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream> {
196196
let options = PyModuleOptions::from_attrs(&mut function.attrs)?;
197-
process_functions_in_module(&options, &mut function)?;
197+
process_functions_in_module(&mut function)?;
198198
let krate = get_pyo3_crate(&options.krate);
199199
let ident = &function.sig.ident;
200200
let vis = &function.vis;
@@ -215,13 +215,13 @@ pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream>
215215
use #krate::impl_::pymodule as impl_;
216216

217217
fn __pyo3_pymodule(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
218-
#ident(module.py(), module.as_gil_ref())
218+
#ident(module.py(), ::std::convert::Into::into(impl_::BoundModule(module)))
219219
}
220220

221221
impl #ident::MakeDef {
222222
const fn make_def() -> impl_::ModuleDef {
223+
const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
223224
unsafe {
224-
const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
225225
impl_::ModuleDef::new(
226226
#ident::__PYO3_NAME,
227227
#doc,
@@ -260,9 +260,8 @@ fn module_initialization(options: PyModuleOptions, ident: &syn::Ident) -> TokenS
260260
}
261261

262262
/// Finds and takes care of the #[pyfn(...)] in `#[pymodule]`
263-
fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn) -> Result<()> {
263+
fn process_functions_in_module(func: &mut syn::ItemFn) -> Result<()> {
264264
let mut stmts: Vec<syn::Stmt> = Vec::new();
265-
let krate = get_pyo3_crate(&options.krate);
266265

267266
for mut stmt in func.block.stmts.drain(..) {
268267
if let syn::Stmt::Item(Item::Fn(func)) = &mut stmt {
@@ -272,7 +271,7 @@ fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn
272271
let name = &func.sig.ident;
273272
let statements: Vec<syn::Stmt> = syn::parse_quote! {
274273
#wrapped_function
275-
#module_name.add_function(#krate::impl_::pyfunction::_wrap_pyfunction(&#name::DEF, #module_name)?)?;
274+
#module_name.add_function(#module_name.wrap_pyfunction(&#name::DEF)?)?;
276275
};
277276
stmts.extend(statements);
278277
}

‎pyo3-macros-backend/src/pyfunction.rs

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -268,12 +268,12 @@ pub fn impl_wrap_pyfunction(
268268
#[doc(hidden)]
269269
#vis mod #name {
270270
pub(crate) struct MakeDef;
271-
pub const DEF: #krate::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
271+
pub const DEF: #krate::impl_::pymethods::PyMethodDef = MakeDef::DEF;
272272

273273
pub fn add_to_module(module: &#krate::Bound<'_, #krate::types::PyModule>) -> #krate::PyResult<()> {
274274
use #krate::prelude::PyModuleMethods;
275275
use ::std::convert::Into;
276-
module.add_function(&#krate::types::PyCFunction::internal_new(&DEF, module.as_gil_ref().into())?)
276+
module.add_function(#krate::types::PyCFunction::internal_new(&DEF, module.as_gil_ref().into())?)
277277
}
278278
}
279279

@@ -284,7 +284,7 @@ pub fn impl_wrap_pyfunction(
284284
const _: () = {
285285
use #krate as _pyo3;
286286
impl #name::MakeDef {
287-
const DEF: #krate::impl_::pyfunction::PyMethodDef = #methoddef;
287+
const DEF: #krate::impl_::pymethods::PyMethodDef = #methoddef;
288288
}
289289

290290
#[allow(non_snake_case)]

‎src/derive_utils.rs

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
//! Functionality for the code generated by the derive backend
22
3-
use crate::{types::PyModule, Python};
3+
use crate::impl_::pymethods::PyMethodDef;
4+
use crate::{
5+
types::{PyCFunction, PyModule},
6+
Bound, PyResult, Python,
7+
};
48

59
/// Enum to abstract over the arguments of Python function wrappers.
610
pub enum PyFunctionArguments<'a> {
@@ -18,6 +22,10 @@ impl<'a> PyFunctionArguments<'a> {
1822
}
1923
}
2024
}
25+
26+
pub fn wrap_pyfunction(self, method_def: &'a PyMethodDef) -> PyResult<&PyCFunction> {
27+
Ok(PyCFunction::internal_new(method_def, self)?.into_gil_ref())
28+
}
2129
}
2230

2331
impl<'a> From<Python<'a>> for PyFunctionArguments<'a> {
@@ -31,3 +39,37 @@ impl<'a> From<&'a PyModule> for PyFunctionArguments<'a> {
3139
PyFunctionArguments::PyModule(module)
3240
}
3341
}
42+
43+
/// Enum to abstract over the arguments of Python function wrappers.
44+
pub enum PyFunctionArgumentsBound<'a, 'py> {
45+
Python(Python<'py>),
46+
PyModule(&'a Bound<'py, PyModule>),
47+
}
48+
49+
impl<'a, 'py> PyFunctionArgumentsBound<'a, 'py> {
50+
pub fn into_py_and_maybe_module(self) -> (Python<'py>, Option<&'a Bound<'py, PyModule>>) {
51+
match self {
52+
PyFunctionArgumentsBound::Python(py) => (py, None),
53+
PyFunctionArgumentsBound::PyModule(module) => {
54+
let py = module.py();
55+
(py, Some(module))
56+
}
57+
}
58+
}
59+
60+
pub fn wrap_pyfunction(self, method_def: &PyMethodDef) -> PyResult<Bound<'py, PyCFunction>> {
61+
PyCFunction::internal_new_bound(method_def, self)
62+
}
63+
}
64+
65+
impl<'a, 'py> From<Python<'py>> for PyFunctionArgumentsBound<'a, 'py> {
66+
fn from(py: Python<'py>) -> PyFunctionArgumentsBound<'a, 'py> {
67+
PyFunctionArgumentsBound::Python(py)
68+
}
69+
}
70+
71+
impl<'a, 'py> From<&'a Bound<'py, PyModule>> for PyFunctionArgumentsBound<'a, 'py> {
72+
fn from(module: &'a Bound<'py, PyModule>) -> PyFunctionArgumentsBound<'a, 'py> {
73+
PyFunctionArgumentsBound::PyModule(module)
74+
}
75+
}

‎src/impl_.rs

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ pub(crate) mod not_send;
1616
pub mod panic;
1717
pub mod pycell;
1818
pub mod pyclass;
19-
pub mod pyfunction;
2019
pub mod pymethods;
2120
pub mod pymodule;
2221
#[doc(hidden)]

‎src/impl_/pyfunction.rs

Lines changed: 0 additions & 10 deletions
This file was deleted.

‎src/impl_/pymodule.rs

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,29 @@ impl<T: PyTypeInfo> PyAddToModule for T {
147147
}
148148
}
149149

150+
pub struct BoundModule<'a, 'py>(pub &'a Bound<'py, PyModule>);
151+
152+
impl<'a> From<BoundModule<'a, 'a>> for &'a PyModule {
153+
#[inline]
154+
fn from(bound: BoundModule<'a, 'a>) -> Self {
155+
bound.0.as_gil_ref()
156+
}
157+
}
158+
159+
impl<'a, 'py> From<BoundModule<'a, 'py>> for &'a Bound<'py, PyModule> {
160+
#[inline]
161+
fn from(bound: BoundModule<'a, 'py>) -> Self {
162+
bound.0
163+
}
164+
}
165+
166+
impl From<BoundModule<'_, '_>> for Py<PyModule> {
167+
#[inline]
168+
fn from(bound: BoundModule<'_, '_>) -> Self {
169+
bound.0.clone().unbind()
170+
}
171+
}
172+
150173
#[cfg(test)]
151174
mod tests {
152175
use std::sync::atomic::{AtomicBool, Ordering};

‎src/macros.rs

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,13 +126,42 @@ macro_rules! py_run_impl {
126126
macro_rules! wrap_pyfunction {
127127
($function:path) => {
128128
&|py_or_module| {
129+
use $crate::derive_utils::PyFunctionArguments;
129130
use $function as wrapped_pyfunction;
130-
$crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, py_or_module)
131+
let function_arguments: PyFunctionArguments<'_> =
132+
::std::convert::Into::into(py_or_module);
133+
function_arguments.wrap_pyfunction(&wrapped_pyfunction::DEF)
131134
}
132135
};
133136
($function:path, $py_or_module:expr) => {{
137+
use $crate::derive_utils::PyFunctionArguments;
134138
use $function as wrapped_pyfunction;
135-
$crate::impl_::pyfunction::_wrap_pyfunction(&wrapped_pyfunction::DEF, $py_or_module)
139+
let function_arguments: PyFunctionArguments<'_> = ::std::convert::Into::into($py_or_module);
140+
function_arguments.wrap_pyfunction(&wrapped_pyfunction::DEF)
141+
}};
142+
}
143+
144+
/// Wraps a Rust function annotated with [`#[pyfunction]`](macro@crate::pyfunction).
145+
///
146+
/// This can be used with [`PyModule::add_function`](crate::types::PyModule::add_function) to add free
147+
/// functions to a [`PyModule`](crate::types::PyModule) - see its documentation for more information.
148+
#[macro_export]
149+
macro_rules! wrap_pyfunction_bound {
150+
($function:path) => {
151+
&|py_or_module| {
152+
use $crate::derive_utils::PyFunctionArgumentsBound;
153+
use $function as wrapped_pyfunction;
154+
let function_arguments: PyFunctionArgumentsBound<'_, '_> =
155+
::std::convert::Into::into($py_or_module);
156+
function_arguments.wrap_pyfunction(&wrapped_pyfunction::DEF)
157+
}
158+
};
159+
($function:path, $py_or_module:expr) => {{
160+
use $crate::derive_utils::PyFunctionArgumentsBound;
161+
use $function as wrapped_pyfunction;
162+
let function_arguments: PyFunctionArgumentsBound<'_, '_> =
163+
::std::convert::Into::into($py_or_module);
164+
function_arguments.wrap_pyfunction(&wrapped_pyfunction::DEF)
136165
}};
137166
}
138167

‎src/prelude.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ pub use crate::PyNativeType;
2323
pub use pyo3_macros::{pyclass, pyfunction, pymethods, pymodule, FromPyObject};
2424

2525
#[cfg(feature = "macros")]
26-
pub use crate::wrap_pyfunction;
26+
pub use crate::{wrap_pyfunction, wrap_pyfunction_bound};
2727

2828
pub use crate::types::any::PyAnyMethods;
2929
pub use crate::types::boolobject::PyBoolMethods;

‎src/types/function.rs

Lines changed: 53 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
1-
use crate::derive_utils::PyFunctionArguments;
1+
use crate::derive_utils::{PyFunctionArguments, PyFunctionArgumentsBound};
22
use crate::ffi_ptr_ext::FfiPtrExt;
33
use crate::methods::PyMethodDefDestructor;
44
use crate::py_result_ext::PyResultExt;
55
use crate::types::capsule::PyCapsuleMethods;
6+
use crate::types::module::PyModuleMethods;
67
use crate::{
78
ffi,
89
impl_::pymethods::{self, PyMethodDef},
@@ -33,17 +34,25 @@ impl PyCFunction {
3334
doc: &'static str,
3435
py_or_module: PyFunctionArguments<'a>,
3536
) -> PyResult<&'a Self> {
36-
Self::new_with_keywords_bound(fun, name, doc, py_or_module).map(Bound::into_gil_ref)
37+
Self::internal_new(
38+
&PyMethodDef::cfunction_with_keywords(
39+
name,
40+
pymethods::PyCFunctionWithKeywords(fun),
41+
doc,
42+
),
43+
py_or_module,
44+
)
45+
.map(Bound::into_gil_ref)
3746
}
3847

3948
/// Create a new built-in function with keywords (*args and/or **kwargs).
40-
pub fn new_with_keywords_bound<'a>(
49+
pub fn new_with_keywords_bound<'py>(
4150
fun: ffi::PyCFunctionWithKeywords,
4251
name: &'static str,
4352
doc: &'static str,
44-
py_or_module: PyFunctionArguments<'a>,
45-
) -> PyResult<Bound<'a, Self>> {
46-
Self::internal_new(
53+
py_or_module: PyFunctionArgumentsBound<'_, 'py>,
54+
) -> PyResult<Bound<'py, Self>> {
55+
Self::internal_new_bound(
4756
&PyMethodDef::cfunction_with_keywords(
4857
name,
4958
pymethods::PyCFunctionWithKeywords(fun),
@@ -67,17 +76,21 @@ impl PyCFunction {
6776
doc: &'static str,
6877
py_or_module: PyFunctionArguments<'a>,
6978
) -> PyResult<&'a Self> {
70-
Self::new_bound(fun, name, doc, py_or_module).map(Bound::into_gil_ref)
79+
Self::internal_new(
80+
&PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
81+
py_or_module,
82+
)
83+
.map(Bound::into_gil_ref)
7184
}
7285

7386
/// Create a new built-in function which takes no arguments.
74-
pub fn new_bound<'a>(
87+
pub fn new_bound<'py>(
7588
fun: ffi::PyCFunction,
7689
name: &'static str,
7790
doc: &'static str,
78-
py_or_module: PyFunctionArguments<'a>,
79-
) -> PyResult<Bound<'a, Self>> {
80-
Self::internal_new(
91+
py_or_module: PyFunctionArgumentsBound<'_, 'py>,
92+
) -> PyResult<Bound<'py, Self>> {
93+
Self::internal_new_bound(
8194
&PyMethodDef::noargs(name, pymethods::PyCFunction(fun), doc),
8295
py_or_module,
8396
)
@@ -189,6 +202,35 @@ impl PyCFunction {
189202
.downcast_into_unchecked()
190203
}
191204
}
205+
206+
#[doc(hidden)]
207+
pub(crate) fn internal_new_bound<'py>(
208+
method_def: &PyMethodDef,
209+
py_or_module: PyFunctionArgumentsBound<'_, 'py>,
210+
) -> PyResult<Bound<'py, Self>> {
211+
let (py, module) = py_or_module.into_py_and_maybe_module();
212+
let (mod_ptr, module_name): (_, Option<Py<PyString>>) = if let Some(m) = module {
213+
let mod_ptr = m.as_ptr();
214+
(mod_ptr, Some(m.name()?.into_py(py)))
215+
} else {
216+
(std::ptr::null_mut(), None)
217+
};
218+
let (def, destructor) = method_def.as_method_def()?;
219+
220+
// FIXME: stop leaking the def and destructor
221+
let def = Box::into_raw(Box::new(def));
222+
std::mem::forget(destructor);
223+
224+
let module_name_ptr = module_name
225+
.as_ref()
226+
.map_or(std::ptr::null_mut(), Py::as_ptr);
227+
228+
unsafe {
229+
ffi::PyCFunction_NewEx(def, mod_ptr, module_name_ptr)
230+
.assume_owned_or_err(py)
231+
.downcast_into_unchecked()
232+
}
233+
}
192234
}
193235

194236
fn closure_capsule_name() -> &'static CStr {

0 commit comments

Comments
 (0)