Skip to content

Commit 6adfe20

Browse files
Tptdavidhewittbirkenfeld
committed
#[pymodule] mod some_module { ... } v3
Based on #2367 and #3294 Allows to export classes, native classes, functions and submodules and provide an init function See test/test_module.rs for an example Future work: - update examples, README and guide - investigate having #[pyclass] and #[pyfunction] directly in the #[pymodule] Co-authored-by: David Hewitt <[email protected]> Co-authored-by: Georg Brandl <[email protected]>
1 parent 9bb0011 commit 6adfe20

20 files changed

+403
-56
lines changed

‎newsfragments/3815.added.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
The ability to create Python modules with a Rust `mod` block.

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@ mod pymethod;
2222
mod quotes;
2323

2424
pub use frompyobject::build_derive_from_pyobject;
25-
pub use module::{process_functions_in_module, pymodule_impl, PyModuleOptions};
25+
pub use module::{
26+
process_functions_in_module, pymodule_function_impl, pymodule_module_impl, PyModuleOptions,
27+
};
2628
pub use pyclass::{build_py_class, build_py_enum, PyClassArgs};
2729
pub use pyfunction::{build_py_function, PyFunctionOptions};
2830
pub use pyimpl::{build_py_methods, PyClassMethodsType};

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

Lines changed: 156 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
33
use crate::{
44
attributes::{self, take_attributes, take_pyo3_options, CrateAttribute, NameAttribute},
5+
get_doc,
56
pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
67
utils::{get_pyo3_crate, PythonDoc},
78
};
@@ -56,9 +57,156 @@ impl PyModuleOptions {
5657
}
5758
}
5859

60+
pub fn pymodule_module_impl(mut module: syn::ItemMod) -> Result<TokenStream> {
61+
let syn::ItemMod {
62+
attrs,
63+
vis,
64+
unsafety: _,
65+
ident,
66+
mod_token,
67+
content,
68+
semi: _,
69+
} = &mut module;
70+
let items = if let Some((_, items)) = content {
71+
items
72+
} else {
73+
bail_spanned!(module.span() => "`#[pymodule]` can only be used on inline modules")
74+
};
75+
let options = PyModuleOptions::from_attrs(attrs)?;
76+
let doc = get_doc(attrs, None);
77+
78+
let name = options.name.unwrap_or_else(|| ident.unraw());
79+
let krate = get_pyo3_crate(&options.krate);
80+
let pyinit_symbol = format!("PyInit_{}", name);
81+
82+
let mut module_items = Vec::new();
83+
let mut module_items_cfg_attrs = Vec::new();
84+
85+
fn extract_use_items(
86+
source: &syn::UseTree,
87+
cfg_attrs: &[syn::Attribute],
88+
target_items: &mut Vec<syn::Ident>,
89+
target_cfg_attrs: &mut Vec<Vec<syn::Attribute>>,
90+
) -> Result<()> {
91+
match source {
92+
syn::UseTree::Name(name) => {
93+
target_items.push(name.ident.clone());
94+
target_cfg_attrs.push(cfg_attrs.to_vec());
95+
}
96+
syn::UseTree::Path(path) => {
97+
extract_use_items(&path.tree, cfg_attrs, target_items, target_cfg_attrs)?
98+
}
99+
syn::UseTree::Group(group) => {
100+
for tree in &group.items {
101+
extract_use_items(tree, cfg_attrs, target_items, target_cfg_attrs)?
102+
}
103+
}
104+
syn::UseTree::Glob(glob) => {
105+
bail_spanned!(glob.span() => "#[pymodule] cannot import glob statements")
106+
}
107+
syn::UseTree::Rename(rename) => {
108+
target_items.push(rename.rename.clone());
109+
target_cfg_attrs.push(cfg_attrs.to_vec());
110+
}
111+
}
112+
Ok(())
113+
}
114+
115+
let mut pymodule_init = None;
116+
117+
for item in &mut *items {
118+
match item {
119+
syn::Item::Use(item_use) => {
120+
let mut is_pyo3 = false;
121+
item_use.attrs.retain(|attr| {
122+
let found = attr.path().is_ident("pyo3");
123+
is_pyo3 |= found;
124+
!found
125+
});
126+
if is_pyo3 {
127+
let cfg_attrs = item_use
128+
.attrs
129+
.iter()
130+
.filter(|attr| attr.path().is_ident("cfg"))
131+
.cloned()
132+
.collect::<Vec<_>>();
133+
extract_use_items(
134+
&item_use.tree,
135+
&cfg_attrs,
136+
&mut module_items,
137+
&mut module_items_cfg_attrs,
138+
)?;
139+
}
140+
}
141+
syn::Item::Fn(item_fn) => {
142+
let mut is_module_init = false;
143+
item_fn.attrs.retain(|attr| {
144+
let found = attr.path().is_ident("pymodule_init");
145+
is_module_init |= found;
146+
!found
147+
});
148+
if is_module_init {
149+
ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one pymodule_init may be specified");
150+
let ident = &item_fn.sig.ident;
151+
pymodule_init = Some(quote! { #ident(module)?; });
152+
}
153+
}
154+
item => {
155+
bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
156+
}
157+
}
158+
}
159+
160+
Ok(quote! {
161+
#vis #mod_token #ident {
162+
#(#items)*
163+
164+
pub const __PYO3_NAME: &'static str = concat!(stringify!(#name), "\0");
165+
166+
pub(crate) struct MakeDef;
167+
impl MakeDef {
168+
const fn make_def() -> #krate::impl_::pymodule::ModuleDef {
169+
use #krate::impl_::pymodule as impl_;
170+
171+
const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
172+
unsafe {
173+
impl_::ModuleDef::new(__PYO3_NAME, #doc, INITIALIZER)
174+
}
175+
}
176+
}
177+
178+
pub static DEF: #krate::impl_::pymodule::ModuleDef = unsafe {
179+
use #krate::impl_::pymodule as impl_;
180+
impl_::ModuleDef::new(concat!(stringify!(#name), "\0"), #doc, impl_::ModuleInitializer(__pyo3_pymodule))
181+
};
182+
183+
pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
184+
module.add_submodule(DEF.make_module(module.py())?.into_ref(module.py()))
185+
}
186+
187+
pub fn __pyo3_pymodule(_py: #krate::Python, module: &#krate::types::PyModule) -> #krate::PyResult<()> {
188+
use #krate::impl_::pymodule::PyAddToModule;
189+
#(
190+
#(#module_items_cfg_attrs)*
191+
#module_items::add_to_module(module)?;
192+
)*
193+
#pymodule_init
194+
Ok(())
195+
}
196+
197+
/// This autogenerated function is called by the python interpreter when importing
198+
/// the module.
199+
#[export_name = #pyinit_symbol]
200+
pub unsafe extern "C" fn __pyo3_init() -> *mut #krate::ffi::PyObject {
201+
#krate::impl_::trampoline::module_init(|py| DEF.make_module(py))
202+
}
203+
}
204+
})
205+
}
206+
59207
/// Generates the function that is called by the python interpreter to initialize the native
60208
/// module
61-
pub fn pymodule_impl(
209+
pub fn pymodule_function_impl(
62210
fnname: &Ident,
63211
options: PyModuleOptions,
64212
doc: PythonDoc,
@@ -75,14 +223,18 @@ pub fn pymodule_impl(
75223
#visibility mod #fnname {
76224
pub(crate) struct MakeDef;
77225
pub static DEF: #krate::impl_::pymodule::ModuleDef = MakeDef::make_def();
78-
pub const NAME: &'static str = concat!(stringify!(#name), "\0");
226+
pub const __PYO3_NAME: &'static str = concat!(stringify!(#name), "\0");
79227

80228
/// This autogenerated function is called by the python interpreter when importing
81229
/// the module.
82230
#[export_name = #pyinit_symbol]
83-
pub unsafe extern "C" fn init() -> *mut #krate::ffi::PyObject {
231+
pub unsafe extern "C" fn __pyo3_init() -> *mut #krate::ffi::PyObject {
84232
#krate::impl_::trampoline::module_init(|py| DEF.make_module(py))
85233
}
234+
235+
pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
236+
module.add_submodule(DEF.make_module(module.py())?.into_ref(module.py()))
237+
}
86238
}
87239

88240
// Generate the definition inside an anonymous function in the same scope as the original function -
@@ -95,7 +247,7 @@ pub fn pymodule_impl(
95247
const fn make_def() -> impl_::ModuleDef {
96248
const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(#fnname);
97249
unsafe {
98-
impl_::ModuleDef::new(#fnname::NAME, #doc, INITIALIZER)
250+
impl_::ModuleDef::new(#fnname::__PYO3_NAME, #doc, INITIALIZER)
99251
}
100252
}
101253
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,10 @@ pub fn impl_wrap_pyfunction(
269269
#vis mod #name {
270270
pub(crate) struct MakeDef;
271271
pub const DEF: #krate::impl_::pyfunction::PyMethodDef = MakeDef::DEF;
272+
273+
pub fn add_to_module(module: &#krate::types::PyModule) -> #krate::PyResult<()> {
274+
module.add_function(#krate::impl_::pyfunction::_wrap_pyfunction(&DEF, module)?)
275+
}
272276
}
273277

274278
// Generate the definition inside an anonymous function in the same scope as the original function -

‎pyo3-macros/src/lib.rs

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@ use proc_macro::TokenStream;
66
use proc_macro2::TokenStream as TokenStream2;
77
use pyo3_macros_backend::{
88
build_derive_from_pyobject, build_py_class, build_py_enum, build_py_function, build_py_methods,
9-
get_doc, process_functions_in_module, pymodule_impl, PyClassArgs, PyClassMethodsType,
10-
PyFunctionOptions, PyModuleOptions,
9+
get_doc, process_functions_in_module, pymodule_function_impl, pymodule_module_impl,
10+
PyClassArgs, PyClassMethodsType, PyFunctionOptions, PyModuleOptions,
1111
};
1212
use quote::quote;
1313
use syn::{parse::Nothing, parse_macro_input};
@@ -37,6 +37,12 @@ use syn::{parse::Nothing, parse_macro_input};
3737
pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
3838
parse_macro_input!(args as Nothing);
3939

40+
if let Ok(module) = syn::parse(input.clone()) {
41+
return pymodule_module_impl(module)
42+
.unwrap_or_compile_error()
43+
.into();
44+
}
45+
4046
let mut ast = parse_macro_input!(input as syn::ItemFn);
4147
let options = match PyModuleOptions::from_attrs(&mut ast.attrs) {
4248
Ok(options) => options,
@@ -49,7 +55,7 @@ pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
4955

5056
let doc = get_doc(&ast.attrs, None);
5157

52-
let expanded = pymodule_impl(&ast.sig.ident, options, doc, &ast.vis);
58+
let expanded = pymodule_function_impl(&ast.sig.ident, options, doc, &ast.vis);
5359

5460
quote!(
5561
#ast

‎pytests/src/lib.rs

Lines changed: 53 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,4 @@
11
use pyo3::prelude::*;
2-
use pyo3::types::PyDict;
3-
use pyo3::wrap_pymodule;
4-
52
pub mod awaitable;
63
pub mod buf_and_str;
74
pub mod comparisons;
@@ -18,43 +15,60 @@ pub mod sequence;
1815
pub mod subclassing;
1916

2017
#[pymodule]
21-
fn pyo3_pytests(py: Python<'_>, m: &PyModule) -> PyResult<()> {
22-
m.add_wrapped(wrap_pymodule!(awaitable::awaitable))?;
18+
mod pyo3_pytests {
19+
use super::*;
20+
#[pyo3]
21+
use awaitable::awaitable;
22+
#[pyo3]
2323
#[cfg(not(Py_LIMITED_API))]
24-
m.add_wrapped(wrap_pymodule!(buf_and_str::buf_and_str))?;
25-
m.add_wrapped(wrap_pymodule!(comparisons::comparisons))?;
24+
use buf_and_str::buf_and_str;
25+
#[pyo3]
26+
use comparisons::comparisons;
2627
#[cfg(not(Py_LIMITED_API))]
27-
m.add_wrapped(wrap_pymodule!(datetime::datetime))?;
28-
m.add_wrapped(wrap_pymodule!(dict_iter::dict_iter))?;
29-
m.add_wrapped(wrap_pymodule!(enums::enums))?;
30-
m.add_wrapped(wrap_pymodule!(misc::misc))?;
31-
m.add_wrapped(wrap_pymodule!(objstore::objstore))?;
32-
m.add_wrapped(wrap_pymodule!(othermod::othermod))?;
33-
m.add_wrapped(wrap_pymodule!(path::path))?;
34-
m.add_wrapped(wrap_pymodule!(pyclasses::pyclasses))?;
35-
m.add_wrapped(wrap_pymodule!(pyfunctions::pyfunctions))?;
36-
m.add_wrapped(wrap_pymodule!(sequence::sequence))?;
37-
m.add_wrapped(wrap_pymodule!(subclassing::subclassing))?;
38-
39-
// Inserting to sys.modules allows importing submodules nicely from Python
40-
// e.g. import pyo3_pytests.buf_and_str as bas
41-
42-
let sys = PyModule::import(py, "sys")?;
43-
let sys_modules: &PyDict = sys.getattr("modules")?.downcast()?;
44-
sys_modules.set_item("pyo3_pytests.awaitable", m.getattr("awaitable")?)?;
45-
sys_modules.set_item("pyo3_pytests.buf_and_str", m.getattr("buf_and_str")?)?;
46-
sys_modules.set_item("pyo3_pytests.comparisons", m.getattr("comparisons")?)?;
47-
sys_modules.set_item("pyo3_pytests.datetime", m.getattr("datetime")?)?;
48-
sys_modules.set_item("pyo3_pytests.dict_iter", m.getattr("dict_iter")?)?;
49-
sys_modules.set_item("pyo3_pytests.enums", m.getattr("enums")?)?;
50-
sys_modules.set_item("pyo3_pytests.misc", m.getattr("misc")?)?;
51-
sys_modules.set_item("pyo3_pytests.objstore", m.getattr("objstore")?)?;
52-
sys_modules.set_item("pyo3_pytests.othermod", m.getattr("othermod")?)?;
53-
sys_modules.set_item("pyo3_pytests.path", m.getattr("path")?)?;
54-
sys_modules.set_item("pyo3_pytests.pyclasses", m.getattr("pyclasses")?)?;
55-
sys_modules.set_item("pyo3_pytests.pyfunctions", m.getattr("pyfunctions")?)?;
56-
sys_modules.set_item("pyo3_pytests.sequence", m.getattr("sequence")?)?;
57-
sys_modules.set_item("pyo3_pytests.subclassing", m.getattr("subclassing")?)?;
28+
#[pyo3]
29+
use datetime::datetime;
30+
#[pyo3]
31+
use dict_iter::dict_iter;
32+
#[pyo3]
33+
use enums::enums;
34+
#[pyo3]
35+
use misc::misc;
36+
#[pyo3]
37+
use objstore::objstore;
38+
#[pyo3]
39+
use othermod::othermod;
40+
#[pyo3]
41+
use path::path;
42+
#[pyo3]
43+
use pyclasses::pyclasses;
44+
#[pyo3]
45+
use pyfunctions::pyfunctions;
46+
use pyo3::types::PyDict;
47+
#[pyo3]
48+
use sequence::sequence;
49+
#[pyo3]
50+
use subclassing::subclassing;
5851

59-
Ok(())
52+
#[pymodule_init]
53+
fn init(m: &PyModule) -> PyResult<()> {
54+
// Inserting to sys.modules allows importing submodules nicely from Python
55+
// e.g. import pyo3_pytests.buf_and_str as bas
56+
let sys = PyModule::import(m.py(), "sys")?;
57+
let sys_modules: &PyDict = sys.getattr("modules")?.downcast()?;
58+
sys_modules.set_item("pyo3_pytests.awaitable", m.getattr("awaitable")?)?;
59+
sys_modules.set_item("pyo3_pytests.buf_and_str", m.getattr("buf_and_str")?)?;
60+
sys_modules.set_item("pyo3_pytests.comparisons", m.getattr("comparisons")?)?;
61+
sys_modules.set_item("pyo3_pytests.datetime", m.getattr("datetime")?)?;
62+
sys_modules.set_item("pyo3_pytests.dict_iter", m.getattr("dict_iter")?)?;
63+
sys_modules.set_item("pyo3_pytests.enums", m.getattr("enums")?)?;
64+
sys_modules.set_item("pyo3_pytests.misc", m.getattr("misc")?)?;
65+
sys_modules.set_item("pyo3_pytests.objstore", m.getattr("objstore")?)?;
66+
sys_modules.set_item("pyo3_pytests.othermod", m.getattr("othermod")?)?;
67+
sys_modules.set_item("pyo3_pytests.path", m.getattr("path")?)?;
68+
sys_modules.set_item("pyo3_pytests.pyclasses", m.getattr("pyclasses")?)?;
69+
sys_modules.set_item("pyo3_pytests.pyfunctions", m.getattr("pyfunctions")?)?;
70+
sys_modules.set_item("pyo3_pytests.sequence", m.getattr("sequence")?)?;
71+
sys_modules.set_item("pyo3_pytests.subclassing", m.getattr("subclassing")?)?;
72+
Ok(())
73+
}
6074
}

‎src/impl_/pymodule.rs

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ use portable_atomic::{AtomicI64, Ordering};
77

88
#[cfg(not(PyPy))]
99
use crate::exceptions::PyImportError;
10-
use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, Python};
10+
use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, PyTypeInfo, Python};
1111

1212
/// `Sync` wrapper of `ffi::PyModuleDef`.
1313
pub struct ModuleDef {
@@ -132,6 +132,17 @@ impl ModuleDef {
132132
}
133133
}
134134

135+
/// Trait to add an element (class, function...) to a module
136+
pub trait PyAddToModule {
137+
fn add_to_module(module: &PyModule) -> PyResult<()>;
138+
}
139+
140+
impl<T: PyTypeInfo> PyAddToModule for T {
141+
fn add_to_module(module: &PyModule) -> PyResult<()> {
142+
module.add(Self::NAME, Self::type_object(module.py()))
143+
}
144+
}
145+
135146
#[cfg(test)]
136147
mod tests {
137148
use std::sync::atomic::{AtomicBool, Ordering};

‎src/macros.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -167,8 +167,8 @@ macro_rules! append_to_inittab {
167167
);
168168
}
169169
$crate::ffi::PyImport_AppendInittab(
170-
$module::NAME.as_ptr() as *const ::std::os::raw::c_char,
171-
::std::option::Option::Some($module::init),
170+
$module::__PYO3_NAME.as_ptr() as *const ::std::os::raw::c_char,
171+
::std::option::Option::Some($module::__pyo3_init),
172172
);
173173
}
174174
};

0 commit comments

Comments
 (0)