1#[cfg(feature = "experimental-inspect")]
4use crate::introspection::{introspection_id_const, module_introspection_code};
5use crate::{
6 attributes::{
7 self, kw, take_attributes, take_pyo3_options, CrateAttribute, GILUsedAttribute,
8 ModuleAttribute, NameAttribute, SubmoduleAttribute,
9 },
10 get_doc,
11 pyclass::PyClassPyO3Option,
12 pyfunction::{impl_wrap_pyfunction, PyFunctionOptions},
13 utils::{has_attribute, has_attribute_with_namespace, Ctx, IdentOrStr, LitCStr},
14};
15use proc_macro2::{Span, TokenStream};
16use quote::quote;
17use std::ffi::CString;
18use syn::{
19 ext::IdentExt,
20 parse::{Parse, ParseStream},
21 parse_quote, parse_quote_spanned,
22 punctuated::Punctuated,
23 spanned::Spanned,
24 token::Comma,
25 Item, Meta, Path, Result,
26};
27
28#[derive(Default)]
29pub struct PyModuleOptions {
30 krate: Option<CrateAttribute>,
31 name: Option<NameAttribute>,
32 module: Option<ModuleAttribute>,
33 submodule: Option<kw::submodule>,
34 gil_used: Option<GILUsedAttribute>,
35}
36
37impl Parse for PyModuleOptions {
38 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
39 let mut options: PyModuleOptions = Default::default();
40
41 options.add_attributes(
42 Punctuated::<PyModulePyO3Option, syn::Token![,]>::parse_terminated(input)?,
43 )?;
44
45 Ok(options)
46 }
47}
48
49impl PyModuleOptions {
50 fn take_pyo3_options(&mut self, attrs: &mut Vec<syn::Attribute>) -> Result<()> {
51 self.add_attributes(take_pyo3_options(attrs)?)
52 }
53
54 fn add_attributes(
55 &mut self,
56 attrs: impl IntoIterator<Item = PyModulePyO3Option>,
57 ) -> Result<()> {
58 macro_rules! set_option {
59 ($key:ident $(, $extra:literal)?) => {
60 {
61 ensure_spanned!(
62 self.$key.is_none(),
63 $key.span() => concat!("`", stringify!($key), "` may only be specified once" $(, $extra)?)
64 );
65 self.$key = Some($key);
66 }
67 };
68 }
69 for attr in attrs {
70 match attr {
71 PyModulePyO3Option::Crate(krate) => set_option!(krate),
72 PyModulePyO3Option::Name(name) => set_option!(name),
73 PyModulePyO3Option::Module(module) => set_option!(module),
74 PyModulePyO3Option::Submodule(submodule) => set_option!(
75 submodule,
76 " (it is implicitly always specified for nested modules)"
77 ),
78 PyModulePyO3Option::GILUsed(gil_used) => {
79 set_option!(gil_used)
80 }
81 }
82 }
83 Ok(())
84 }
85}
86
87pub fn pymodule_module_impl(
88 module: &mut syn::ItemMod,
89 mut options: PyModuleOptions,
90) -> Result<TokenStream> {
91 let syn::ItemMod {
92 attrs,
93 vis,
94 unsafety: _,
95 ident,
96 mod_token,
97 content,
98 semi: _,
99 } = module;
100 let items = if let Some((_, items)) = content {
101 items
102 } else {
103 bail_spanned!(mod_token.span() => "`#[pymodule]` can only be used on inline modules")
104 };
105 options.take_pyo3_options(attrs)?;
106 let ctx = &Ctx::new(&options.krate, None);
107 let Ctx { pyo3_path, .. } = ctx;
108 let doc = get_doc(attrs, None, ctx);
109 let name = options
110 .name
111 .map_or_else(|| ident.unraw(), |name| name.value.0);
112 let full_name = if let Some(module) = &options.module {
113 format!("{}.{}", module.value.value(), name)
114 } else {
115 name.to_string()
116 };
117
118 let mut module_items = Vec::new();
119 let mut module_items_cfg_attrs = Vec::new();
120
121 fn extract_use_items(
122 source: &syn::UseTree,
123 cfg_attrs: &[syn::Attribute],
124 target_items: &mut Vec<syn::Ident>,
125 target_cfg_attrs: &mut Vec<Vec<syn::Attribute>>,
126 ) -> Result<()> {
127 match source {
128 syn::UseTree::Name(name) => {
129 target_items.push(name.ident.clone());
130 target_cfg_attrs.push(cfg_attrs.to_vec());
131 }
132 syn::UseTree::Path(path) => {
133 extract_use_items(&path.tree, cfg_attrs, target_items, target_cfg_attrs)?
134 }
135 syn::UseTree::Group(group) => {
136 for tree in &group.items {
137 extract_use_items(tree, cfg_attrs, target_items, target_cfg_attrs)?
138 }
139 }
140 syn::UseTree::Glob(glob) => {
141 bail_spanned!(glob.span() => "#[pymodule] cannot import glob statements")
142 }
143 syn::UseTree::Rename(rename) => {
144 target_items.push(rename.rename.clone());
145 target_cfg_attrs.push(cfg_attrs.to_vec());
146 }
147 }
148 Ok(())
149 }
150
151 let mut pymodule_init = None;
152 let mut module_consts = Vec::new();
153 let mut module_consts_cfg_attrs = Vec::new();
154
155 for item in &mut *items {
156 match item {
157 Item::Use(item_use) => {
158 let is_pymodule_export =
159 find_and_remove_attribute(&mut item_use.attrs, "pymodule_export");
160 if is_pymodule_export {
161 let cfg_attrs = get_cfg_attributes(&item_use.attrs);
162 extract_use_items(
163 &item_use.tree,
164 &cfg_attrs,
165 &mut module_items,
166 &mut module_items_cfg_attrs,
167 )?;
168 }
169 }
170 Item::Fn(item_fn) => {
171 ensure_spanned!(
172 !has_attribute(&item_fn.attrs, "pymodule_export"),
173 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
174 );
175 let is_pymodule_init =
176 find_and_remove_attribute(&mut item_fn.attrs, "pymodule_init");
177 let ident = &item_fn.sig.ident;
178 if is_pymodule_init {
179 ensure_spanned!(
180 !has_attribute(&item_fn.attrs, "pyfunction"),
181 item_fn.span() => "`#[pyfunction]` cannot be used alongside `#[pymodule_init]`"
182 );
183 ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one `#[pymodule_init]` may be specified");
184 pymodule_init = Some(quote! { #ident(module)?; });
185 } else if has_attribute(&item_fn.attrs, "pyfunction")
186 || has_attribute_with_namespace(
187 &item_fn.attrs,
188 Some(pyo3_path),
189 &["pyfunction"],
190 )
191 || has_attribute_with_namespace(
192 &item_fn.attrs,
193 Some(pyo3_path),
194 &["prelude", "pyfunction"],
195 )
196 {
197 module_items.push(ident.clone());
198 module_items_cfg_attrs.push(get_cfg_attributes(&item_fn.attrs));
199 }
200 }
201 Item::Struct(item_struct) => {
202 ensure_spanned!(
203 !has_attribute(&item_struct.attrs, "pymodule_export"),
204 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
205 );
206 if has_attribute(&item_struct.attrs, "pyclass")
207 || has_attribute_with_namespace(
208 &item_struct.attrs,
209 Some(pyo3_path),
210 &["pyclass"],
211 )
212 || has_attribute_with_namespace(
213 &item_struct.attrs,
214 Some(pyo3_path),
215 &["prelude", "pyclass"],
216 )
217 {
218 module_items.push(item_struct.ident.clone());
219 module_items_cfg_attrs.push(get_cfg_attributes(&item_struct.attrs));
220 if !has_pyo3_module_declared::<PyClassPyO3Option>(
221 &item_struct.attrs,
222 "pyclass",
223 |option| matches!(option, PyClassPyO3Option::Module(_)),
224 )? {
225 set_module_attribute(&mut item_struct.attrs, &full_name);
226 }
227 }
228 }
229 Item::Enum(item_enum) => {
230 ensure_spanned!(
231 !has_attribute(&item_enum.attrs, "pymodule_export"),
232 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
233 );
234 if has_attribute(&item_enum.attrs, "pyclass")
235 || has_attribute_with_namespace(&item_enum.attrs, Some(pyo3_path), &["pyclass"])
236 || has_attribute_with_namespace(
237 &item_enum.attrs,
238 Some(pyo3_path),
239 &["prelude", "pyclass"],
240 )
241 {
242 module_items.push(item_enum.ident.clone());
243 module_items_cfg_attrs.push(get_cfg_attributes(&item_enum.attrs));
244 if !has_pyo3_module_declared::<PyClassPyO3Option>(
245 &item_enum.attrs,
246 "pyclass",
247 |option| matches!(option, PyClassPyO3Option::Module(_)),
248 )? {
249 set_module_attribute(&mut item_enum.attrs, &full_name);
250 }
251 }
252 }
253 Item::Mod(item_mod) => {
254 ensure_spanned!(
255 !has_attribute(&item_mod.attrs, "pymodule_export"),
256 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
257 );
258 if has_attribute(&item_mod.attrs, "pymodule")
259 || has_attribute_with_namespace(&item_mod.attrs, Some(pyo3_path), &["pymodule"])
260 || has_attribute_with_namespace(
261 &item_mod.attrs,
262 Some(pyo3_path),
263 &["prelude", "pymodule"],
264 )
265 {
266 module_items.push(item_mod.ident.clone());
267 module_items_cfg_attrs.push(get_cfg_attributes(&item_mod.attrs));
268 if !has_pyo3_module_declared::<PyModulePyO3Option>(
269 &item_mod.attrs,
270 "pymodule",
271 |option| matches!(option, PyModulePyO3Option::Module(_)),
272 )? {
273 set_module_attribute(&mut item_mod.attrs, &full_name);
274 }
275 item_mod
276 .attrs
277 .push(parse_quote_spanned!(item_mod.mod_token.span()=> #[pyo3(submodule)]));
278 }
279 }
280 Item::ForeignMod(item) => {
281 ensure_spanned!(
282 !has_attribute(&item.attrs, "pymodule_export"),
283 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
284 );
285 }
286 Item::Trait(item) => {
287 ensure_spanned!(
288 !has_attribute(&item.attrs, "pymodule_export"),
289 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
290 );
291 }
292 Item::Const(item) => {
293 if !find_and_remove_attribute(&mut item.attrs, "pymodule_export") {
294 continue;
295 }
296
297 module_consts.push(item.ident.clone());
298 module_consts_cfg_attrs.push(get_cfg_attributes(&item.attrs));
299 }
300 Item::Static(item) => {
301 ensure_spanned!(
302 !has_attribute(&item.attrs, "pymodule_export"),
303 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
304 );
305 }
306 Item::Macro(item) => {
307 ensure_spanned!(
308 !has_attribute(&item.attrs, "pymodule_export"),
309 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
310 );
311 }
312 Item::ExternCrate(item) => {
313 ensure_spanned!(
314 !has_attribute(&item.attrs, "pymodule_export"),
315 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
316 );
317 }
318 Item::Impl(item) => {
319 ensure_spanned!(
320 !has_attribute(&item.attrs, "pymodule_export"),
321 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
322 );
323 }
324 Item::TraitAlias(item) => {
325 ensure_spanned!(
326 !has_attribute(&item.attrs, "pymodule_export"),
327 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
328 );
329 }
330 Item::Type(item) => {
331 ensure_spanned!(
332 !has_attribute(&item.attrs, "pymodule_export"),
333 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
334 );
335 }
336 Item::Union(item) => {
337 ensure_spanned!(
338 !has_attribute(&item.attrs, "pymodule_export"),
339 item.span() => "`#[pymodule_export]` may only be used on `use` or `const` statements"
340 );
341 }
342 _ => (),
343 }
344 }
345
346 #[cfg(feature = "experimental-inspect")]
347 let introspection = module_introspection_code(
348 pyo3_path,
349 &name.to_string(),
350 &module_items,
351 &module_items_cfg_attrs,
352 );
353 #[cfg(not(feature = "experimental-inspect"))]
354 let introspection = quote! {};
355 #[cfg(feature = "experimental-inspect")]
356 let introspection_id = introspection_id_const();
357 #[cfg(not(feature = "experimental-inspect"))]
358 let introspection_id = quote! {};
359
360 let module_def = quote! {{
361 use #pyo3_path::impl_::pymodule as impl_;
362 const INITIALIZER: impl_::ModuleInitializer = impl_::ModuleInitializer(__pyo3_pymodule);
363 unsafe {
364 impl_::ModuleDef::new(
365 __PYO3_NAME,
366 #doc,
367 INITIALIZER
368 )
369 }
370 }};
371 let initialization = module_initialization(
372 &name,
373 ctx,
374 module_def,
375 options.submodule.is_some(),
376 options.gil_used.map_or(true, |op| op.value.value),
377 );
378
379 let module_consts_names = module_consts.iter().map(|i| i.unraw().to_string());
380
381 Ok(quote!(
382 #(#attrs)*
383 #vis #mod_token #ident {
384 #(#items)*
385
386 #initialization
387 #introspection
388 #introspection_id
389
390 fn __pyo3_pymodule(module: &#pyo3_path::Bound<'_, #pyo3_path::types::PyModule>) -> #pyo3_path::PyResult<()> {
391 use #pyo3_path::impl_::pymodule::PyAddToModule;
392 #(
393 #(#module_items_cfg_attrs)*
394 #module_items::_PYO3_DEF.add_to_module(module)?;
395 )*
396
397 #(
398 #(#module_consts_cfg_attrs)*
399 #pyo3_path::types::PyModuleMethods::add(module, #module_consts_names, #module_consts)?;
400 )*
401
402 #pymodule_init
403 ::std::result::Result::Ok(())
404 }
405 }
406 ))
407}
408
409pub fn pymodule_function_impl(
412 function: &mut syn::ItemFn,
413 mut options: PyModuleOptions,
414) -> Result<TokenStream> {
415 options.take_pyo3_options(&mut function.attrs)?;
416 process_functions_in_module(&options, function)?;
417 let ctx = &Ctx::new(&options.krate, None);
418 let Ctx { pyo3_path, .. } = ctx;
419 let ident = &function.sig.ident;
420 let name = options
421 .name
422 .map_or_else(|| ident.unraw(), |name| name.value.0);
423 let vis = &function.vis;
424 let doc = get_doc(&function.attrs, None, ctx);
425
426 let initialization = module_initialization(
427 &name,
428 ctx,
429 quote! { MakeDef::make_def() },
430 false,
431 options.gil_used.map_or(true, |op| op.value.value),
432 );
433
434 #[cfg(feature = "experimental-inspect")]
435 let introspection = module_introspection_code(pyo3_path, &name.to_string(), &[], &[]);
436 #[cfg(not(feature = "experimental-inspect"))]
437 let introspection = quote! {};
438 #[cfg(feature = "experimental-inspect")]
439 let introspection_id = introspection_id_const();
440 #[cfg(not(feature = "experimental-inspect"))]
441 let introspection_id = quote! {};
442
443 let mut module_args = Vec::new();
445 if function.sig.inputs.len() == 2 {
446 module_args.push(quote!(module.py()));
447 }
448 module_args
449 .push(quote!(::std::convert::Into::into(#pyo3_path::impl_::pymethods::BoundRef(module))));
450
451 Ok(quote! {
452 #[doc(hidden)]
453 #vis mod #ident {
454 #initialization
455 #introspection
456 #introspection_id
457 }
458
459 #[allow(unknown_lints, non_local_definitions)]
464 impl #ident::MakeDef {
465 const fn make_def() -> #pyo3_path::impl_::pymodule::ModuleDef {
466 fn __pyo3_pymodule(module: &#pyo3_path::Bound<'_, #pyo3_path::types::PyModule>) -> #pyo3_path::PyResult<()> {
467 #ident(#(#module_args),*)
468 }
469
470 const INITIALIZER: #pyo3_path::impl_::pymodule::ModuleInitializer = #pyo3_path::impl_::pymodule::ModuleInitializer(__pyo3_pymodule);
471 unsafe {
472 #pyo3_path::impl_::pymodule::ModuleDef::new(
473 #ident::__PYO3_NAME,
474 #doc,
475 INITIALIZER
476 )
477 }
478 }
479 }
480 })
481}
482
483fn module_initialization(
484 name: &syn::Ident,
485 ctx: &Ctx,
486 module_def: TokenStream,
487 is_submodule: bool,
488 gil_used: bool,
489) -> TokenStream {
490 let Ctx { pyo3_path, .. } = ctx;
491 let pyinit_symbol = format!("PyInit_{name}");
492 let name = name.to_string();
493 let pyo3_name = LitCStr::new(CString::new(name).unwrap(), Span::call_site(), ctx);
494
495 let mut result = quote! {
496 #[doc(hidden)]
497 pub const __PYO3_NAME: &'static ::std::ffi::CStr = #pyo3_name;
498
499 pub(super) struct MakeDef;
500 #[doc(hidden)]
501 pub static _PYO3_DEF: #pyo3_path::impl_::pymodule::ModuleDef = #module_def;
502 #[doc(hidden)]
503 pub static __PYO3_GIL_USED: bool = #gil_used;
505 };
506 if !is_submodule {
507 result.extend(quote! {
508 #[doc(hidden)]
511 #[export_name = #pyinit_symbol]
512 pub unsafe extern "C" fn __pyo3_init() -> *mut #pyo3_path::ffi::PyObject {
513 unsafe { #pyo3_path::impl_::trampoline::module_init(|py| _PYO3_DEF.make_module(py, #gil_used)) }
514 }
515 });
516 }
517 result
518}
519
520fn process_functions_in_module(options: &PyModuleOptions, func: &mut syn::ItemFn) -> Result<()> {
522 let ctx = &Ctx::new(&options.krate, None);
523 let Ctx { pyo3_path, .. } = ctx;
524 let mut stmts: Vec<syn::Stmt> = Vec::new();
525
526 for mut stmt in func.block.stmts.drain(..) {
527 if let syn::Stmt::Item(Item::Fn(func)) = &mut stmt {
528 if let Some(pyfn_args) = get_pyfn_attr(&mut func.attrs)? {
529 let module_name = pyfn_args.modname;
530 let wrapped_function = impl_wrap_pyfunction(func, pyfn_args.options)?;
531 let name = &func.sig.ident;
532 let statements: Vec<syn::Stmt> = syn::parse_quote! {
533 #wrapped_function
534 {
535 use #pyo3_path::types::PyModuleMethods;
536 #module_name.add_function(#pyo3_path::wrap_pyfunction!(#name, #module_name.as_borrowed())?)?;
537 }
538 };
539 stmts.extend(statements);
540 }
541 };
542 stmts.push(stmt);
543 }
544
545 func.block.stmts = stmts;
546 Ok(())
547}
548
549pub struct PyFnArgs {
550 modname: Path,
551 options: PyFunctionOptions,
552}
553
554impl Parse for PyFnArgs {
555 fn parse(input: syn::parse::ParseStream<'_>) -> syn::Result<Self> {
556 let modname = input.parse().map_err(
557 |e| err_spanned!(e.span() => "expected module as first argument to #[pyfn()]"),
558 )?;
559
560 if input.is_empty() {
561 return Ok(Self {
562 modname,
563 options: Default::default(),
564 });
565 }
566
567 let _: Comma = input.parse()?;
568
569 Ok(Self {
570 modname,
571 options: input.parse()?,
572 })
573 }
574}
575
576fn get_pyfn_attr(attrs: &mut Vec<syn::Attribute>) -> syn::Result<Option<PyFnArgs>> {
578 let mut pyfn_args: Option<PyFnArgs> = None;
579
580 take_attributes(attrs, |attr| {
581 if attr.path().is_ident("pyfn") {
582 ensure_spanned!(
583 pyfn_args.is_none(),
584 attr.span() => "`#[pyfn] may only be specified once"
585 );
586 pyfn_args = Some(attr.parse_args()?);
587 Ok(true)
588 } else {
589 Ok(false)
590 }
591 })?;
592
593 if let Some(pyfn_args) = &mut pyfn_args {
594 pyfn_args
595 .options
596 .add_attributes(take_pyo3_options(attrs)?)?;
597 }
598
599 Ok(pyfn_args)
600}
601
602fn get_cfg_attributes(attrs: &[syn::Attribute]) -> Vec<syn::Attribute> {
603 attrs
604 .iter()
605 .filter(|attr| attr.path().is_ident("cfg"))
606 .cloned()
607 .collect()
608}
609
610fn find_and_remove_attribute(attrs: &mut Vec<syn::Attribute>, ident: &str) -> bool {
611 let mut found = false;
612 attrs.retain(|attr| {
613 if attr.path().is_ident(ident) {
614 found = true;
615 false
616 } else {
617 true
618 }
619 });
620 found
621}
622
623impl PartialEq<syn::Ident> for IdentOrStr<'_> {
624 fn eq(&self, other: &syn::Ident) -> bool {
625 match self {
626 IdentOrStr::Str(s) => other == s,
627 IdentOrStr::Ident(i) => other == i,
628 }
629 }
630}
631
632fn set_module_attribute(attrs: &mut Vec<syn::Attribute>, module_name: &str) {
633 attrs.push(parse_quote!(#[pyo3(module = #module_name)]));
634}
635
636fn has_pyo3_module_declared<T: Parse>(
637 attrs: &[syn::Attribute],
638 root_attribute_name: &str,
639 is_module_option: impl Fn(&T) -> bool + Copy,
640) -> Result<bool> {
641 for attr in attrs {
642 if (attr.path().is_ident("pyo3") || attr.path().is_ident(root_attribute_name))
643 && matches!(attr.meta, Meta::List(_))
644 {
645 for option in &attr.parse_args_with(Punctuated::<T, Comma>::parse_terminated)? {
646 if is_module_option(option) {
647 return Ok(true);
648 }
649 }
650 }
651 }
652 Ok(false)
653}
654
655enum PyModulePyO3Option {
656 Submodule(SubmoduleAttribute),
657 Crate(CrateAttribute),
658 Name(NameAttribute),
659 Module(ModuleAttribute),
660 GILUsed(GILUsedAttribute),
661}
662
663impl Parse for PyModulePyO3Option {
664 fn parse(input: ParseStream<'_>) -> Result<Self> {
665 let lookahead = input.lookahead1();
666 if lookahead.peek(attributes::kw::name) {
667 input.parse().map(PyModulePyO3Option::Name)
668 } else if lookahead.peek(syn::Token![crate]) {
669 input.parse().map(PyModulePyO3Option::Crate)
670 } else if lookahead.peek(attributes::kw::module) {
671 input.parse().map(PyModulePyO3Option::Module)
672 } else if lookahead.peek(attributes::kw::submodule) {
673 input.parse().map(PyModulePyO3Option::Submodule)
674 } else if lookahead.peek(attributes::kw::gil_used) {
675 input.parse().map(PyModulePyO3Option::GILUsed)
676 } else {
677 Err(lookahead.error())
678 }
679 }
680}