Skip to content

Commit

Permalink
allow #[pymodule(...)] to accept all relevant #[pyo3(...)] options (
Browse files Browse the repository at this point in the history
  • Loading branch information
davidhewitt committed Jul 17, 2024
1 parent 6317570 commit 7e0fc30
Show file tree
Hide file tree
Showing 16 changed files with 108 additions and 107 deletions.
3 changes: 1 addition & 2 deletions examples/getitem/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,7 @@ impl ExampleContainer {
}
}

#[pymodule]
#[pyo3(name = "getitem")]
#[pymodule(name = "getitem")]
fn example(m: &Bound<'_, PyModule>) -> PyResult<()> {
// ? -https://github.com/PyO3/maturin/issues/475
m.add_class::<ExampleContainer>()?;
Expand Down
3 changes: 1 addition & 2 deletions guide/src/module.md
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,7 @@ fn double(x: usize) -> usize {
x * 2
}

#[pymodule]
#[pyo3(name = "custom_name")]
#[pymodule(name = "custom_name")]
fn my_extension(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(double, m)?)
}
Expand Down
1 change: 1 addition & 0 deletions newsfragments/4330.changed.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
`#[pymodule(...)]` now directly accepts all relevant `#[pyo3(...)]` options.
121 changes: 57 additions & 64 deletions pyo3-macros-backend/src/module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
use crate::{
attributes::{
self, take_attributes, take_pyo3_options, CrateAttribute, ModuleAttribute, NameAttribute,
SubmoduleAttribute,
self, kw, take_attributes, take_pyo3_options, CrateAttribute, ModuleAttribute,
NameAttribute, SubmoduleAttribute,
},
get_doc,
pyclass::PyClassPyO3Option,
Expand All @@ -26,97 +26,85 @@ use syn::{
#[derive(Default)]
pub struct PyModuleOptions {
krate: Option<CrateAttribute>,
name: Option<syn::Ident>,
name: Option<NameAttribute>,
module: Option<ModuleAttribute>,
is_submodule: bool,
submodule: Option<kw::submodule>,
}

impl PyModuleOptions {
pub fn from_attrs(attrs: &mut Vec<syn::Attribute>) -> Result<Self> {
impl Parse for PyModuleOptions {
fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
let mut options: PyModuleOptions = Default::default();

for option in take_pyo3_options(attrs)? {
match option {
PyModulePyO3Option::Name(name) => options.set_name(name.value.0)?,
PyModulePyO3Option::Crate(path) => options.set_crate(path)?,
PyModulePyO3Option::Module(module) => options.set_module(module)?,
PyModulePyO3Option::Submodule(submod) => options.set_submodule(submod)?,
}
}
options.add_attributes(
Punctuated::<PyModulePyO3Option, syn::Token![,]>::parse_terminated(input)?,
)?;

Ok(options)
}
}

fn set_name(&mut self, name: syn::Ident) -> Result<()> {
ensure_spanned!(
self.name.is_none(),
name.span() => "`name` may only be specified once"
);

self.name = Some(name);
Ok(())
}

fn set_crate(&mut self, path: CrateAttribute) -> Result<()> {
ensure_spanned!(
self.krate.is_none(),
path.span() => "`crate` may only be specified once"
);

self.krate = Some(path);
Ok(())
}

fn set_module(&mut self, name: ModuleAttribute) -> Result<()> {
ensure_spanned!(
self.module.is_none(),
name.span() => "`module` may only be specified once"
);

self.module = Some(name);
Ok(())
impl PyModuleOptions {
fn take_pyo3_options(&mut self, attrs: &mut Vec<syn::Attribute>) -> Result<()> {
self.add_attributes(take_pyo3_options(attrs)?)
}

fn set_submodule(&mut self, submod: SubmoduleAttribute) -> Result<()> {
ensure_spanned!(
!self.is_submodule,
submod.span() => "`submodule` may only be specified once"
);

self.is_submodule = true;
fn add_attributes(
&mut self,
attrs: impl IntoIterator<Item = PyModulePyO3Option>,
) -> Result<()> {
macro_rules! set_option {
($key:ident) => {
{
ensure_spanned!(
self.$key.is_none(),
$key.span() => concat!("`", stringify!($key), "` may only be specified once")
);
self.$key = Some($key);
}
};
}
for attr in attrs {
match attr {
PyModulePyO3Option::Crate(krate) => set_option!(krate),
PyModulePyO3Option::Name(name) => set_option!(name),
PyModulePyO3Option::Module(module) => set_option!(module),
PyModulePyO3Option::Submodule(submodule) => set_option!(submodule),
}
}
Ok(())
}
}

pub fn pymodule_module_impl(
mut module: syn::ItemMod,
mut is_submodule: bool,
module: &mut syn::ItemMod,
mut options: PyModuleOptions,
) -> Result<TokenStream> {
let syn::ItemMod {
attrs,
vis,
unsafety: _,
ident,
mod_token: _,
mod_token,
content,
semi: _,
} = &mut module;
} = module;
let items = if let Some((_, items)) = content {
items
} else {
bail_spanned!(module.span() => "`#[pymodule]` can only be used on inline modules")
bail_spanned!(mod_token.span() => "`#[pymodule]` can only be used on inline modules")
};
let options = PyModuleOptions::from_attrs(attrs)?;
options.take_pyo3_options(attrs)?;
let ctx = &Ctx::new(&options.krate, None);
let Ctx { pyo3_path, .. } = ctx;
let doc = get_doc(attrs, None, ctx);
let name = options.name.unwrap_or_else(|| ident.unraw());
let name = options
.name
.map_or_else(|| ident.unraw(), |name| name.value.0);
let full_name = if let Some(module) = &options.module {
format!("{}.{}", module.value.value(), name)
} else {
name.to_string()
};
is_submodule = is_submodule || options.is_submodule;

let mut module_items = Vec::new();
let mut module_items_cfg_attrs = Vec::new();
Expand Down Expand Up @@ -350,10 +338,11 @@ pub fn pymodule_module_impl(
)
}
}};
let initialization = module_initialization(&name, ctx, module_def, is_submodule);
let initialization = module_initialization(&name, ctx, module_def, options.submodule.is_some());

Ok(quote!(
#(#attrs)*
#vis mod #ident {
#vis #mod_token #ident {
#(#items)*

#initialization
Expand All @@ -373,14 +362,19 @@ pub fn pymodule_module_impl(

/// Generates the function that is called by the python interpreter to initialize the native
/// module
pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream> {
let options = PyModuleOptions::from_attrs(&mut function.attrs)?;
process_functions_in_module(&options, &mut function)?;
pub fn pymodule_function_impl(
function: &mut syn::ItemFn,
mut options: PyModuleOptions,
) -> Result<TokenStream> {
options.take_pyo3_options(&mut function.attrs)?;
process_functions_in_module(&options, function)?;
let ctx = &Ctx::new(&options.krate, None);
let stmts = std::mem::take(&mut function.block.stmts);
let Ctx { pyo3_path, .. } = ctx;
let ident = &function.sig.ident;
let name = options.name.unwrap_or_else(|| ident.unraw());
let name = options
.name
.map_or_else(|| ident.unraw(), |name| name.value.0);
let vis = &function.vis;
let doc = get_doc(&function.attrs, None, ctx);

Expand Down Expand Up @@ -419,7 +413,6 @@ pub fn pymodule_function_impl(mut function: syn::ItemFn) -> Result<TokenStream>
.push(parse_quote!(#[allow(clippy::used_underscore_binding)]));

Ok(quote! {
#function
#[doc(hidden)]
#vis mod #ident {
#initialization
Expand Down
44 changes: 22 additions & 22 deletions pyo3-macros/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@
#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
use proc_macro::TokenStream;
use proc_macro2::{Span, TokenStream as TokenStream2};
use proc_macro2::TokenStream as TokenStream2;
use pyo3_macros_backend::{
build_derive_from_pyobject, build_py_class, build_py_enum, build_py_function, build_py_methods,
pymodule_function_impl, pymodule_module_impl, PyClassArgs, PyClassMethodsType,
PyFunctionOptions,
PyFunctionOptions, PyModuleOptions,
};
use quote::quote;
use syn::{parse::Nothing, parse_macro_input, Item};
use syn::{parse_macro_input, Item};

/// A proc macro used to implement Python modules.
///
Expand All @@ -24,6 +24,9 @@ use syn::{parse::Nothing, parse_macro_input, Item};
/// | Annotation | Description |
/// | :- | :- |
/// | `#[pyo3(name = "...")]` | Defines the name of the module in Python. |
/// | `#[pyo3(submodule)]` | Skips adding a `PyInit_` FFI symbol to the compiled binary. |
/// | `#[pyo3(module = "...")]` | Defines the Python `dotted.path` to the parent module for use in introspection. |
/// | `#[pyo3(crate = "pyo3")]` | Defines the path to PyO3 to use code generated by the macro. |
///
/// For more on creating Python modules see the [module section of the guide][1].
///
Expand All @@ -35,32 +38,29 @@ use syn::{parse::Nothing, parse_macro_input, Item};
#[doc = concat!("[1]: https://pyo3.rs/v", env!("CARGO_PKG_VERSION"), "/module.html")]
#[proc_macro_attribute]
pub fn pymodule(args: TokenStream, input: TokenStream) -> TokenStream {
match parse_macro_input!(input as Item) {
let options = parse_macro_input!(args as PyModuleOptions);

let mut ast = parse_macro_input!(input as Item);
let expanded = match &mut ast {
Item::Mod(module) => {
let is_submodule = match parse_macro_input!(args as Option<syn::Ident>) {
Some(i) if i == "submodule" => true,
Some(_) => {
return syn::Error::new(
Span::call_site(),
"#[pymodule] only accepts submodule as an argument",
)
.into_compile_error()
.into();
}
None => false,
};
pymodule_module_impl(module, is_submodule)
}
Item::Fn(function) => {
parse_macro_input!(args as Nothing);
pymodule_function_impl(function)
match pymodule_module_impl(module, options) {
// #[pymodule] on a module will rebuild the original ast, so we don't emit it here
Ok(expanded) => return expanded.into(),
Err(e) => Err(e),
}
}
Item::Fn(function) => pymodule_function_impl(function, options),
unsupported => Err(syn::Error::new_spanned(
unsupported,
"#[pymodule] only supports modules and functions.",
)),
}
.unwrap_or_compile_error()
.unwrap_or_compile_error();

quote!(
#ast
#expanded
)
.into()
}

Expand Down
6 changes: 2 additions & 4 deletions tests/test_declarative_module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,7 @@ create_exception!(
"Some description."
);

#[pymodule]
#[pyo3(submodule)]
#[pymodule(submodule)]
mod external_submodule {}

/// A module written using declarative syntax.
Expand Down Expand Up @@ -144,8 +143,7 @@ mod declarative_submodule {
use super::{double, double_value};
}

#[pymodule]
#[pyo3(name = "declarative_module_renamed")]
#[pymodule(name = "declarative_module_renamed")]
mod declarative_module2 {
#[pymodule_export]
use super::double;
Expand Down
3 changes: 1 addition & 2 deletions tests/test_module.rs
Original file line number Diff line number Diff line change
Expand Up @@ -138,8 +138,7 @@ fn test_module_with_explicit_py_arg() {
});
}

#[pymodule]
#[pyo3(name = "other_name")]
#[pymodule(name = "other_name")]
fn some_name(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add("other_name", "other_name")?;
Ok(())
Expand Down
1 change: 1 addition & 0 deletions tests/ui/empty.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
// see invalid_pymodule_in_root.rs
2 changes: 1 addition & 1 deletion tests/ui/invalid_pymodule_args.stderr
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
error: unexpected token
error: expected one of: `name`, `crate`, `module`, `submodule`
--> tests/ui/invalid_pymodule_args.rs:3:12
|
3 | #[pymodule(some_arg)]
Expand Down
2 changes: 2 additions & 0 deletions tests/ui/invalid_pymodule_glob.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
#![allow(unused_imports)]

use pyo3::prelude::*;

#[pyfunction]
Expand Down
4 changes: 2 additions & 2 deletions tests/ui/invalid_pymodule_glob.stderr
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
error: #[pymodule] cannot import glob statements
--> tests/ui/invalid_pymodule_glob.rs:11:16
--> tests/ui/invalid_pymodule_glob.rs:13:16
|
11 | use super::*;
13 | use super::*;
| ^
1 change: 1 addition & 0 deletions tests/ui/invalid_pymodule_in_root.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use pyo3::prelude::*;

#[pymodule]
#[path = "empty.rs"] // to silence error related to missing file
mod invalid_pymodule_in_root_module;

fn main() {}
8 changes: 4 additions & 4 deletions tests/ui/invalid_pymodule_in_root.stderr
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
error[E0658]: non-inline modules in proc macro input are unstable
--> tests/ui/invalid_pymodule_in_root.rs:4:1
--> tests/ui/invalid_pymodule_in_root.rs:5:1
|
4 | mod invalid_pymodule_in_root_module;
5 | mod invalid_pymodule_in_root_module;
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
|
= note: see issue #54727 <https://github.com/rust-lang/rust/issues/54727> for more information

error: `#[pymodule]` can only be used on inline modules
--> tests/ui/invalid_pymodule_in_root.rs:4:1
--> tests/ui/invalid_pymodule_in_root.rs:5:1
|
4 | mod invalid_pymodule_in_root_module;
5 | mod invalid_pymodule_in_root_module;
| ^^^
6 changes: 6 additions & 0 deletions tests/ui/invalid_pymodule_trait.stderr
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,9 @@ error: `#[pymodule_export]` may only be used on `use` statements
|
5 | #[pymodule_export]
| ^

error: cannot find attribute `pymodule_export` in this scope
--> tests/ui/invalid_pymodule_trait.rs:5:7
|
5 | #[pymodule_export]
| ^^^^^^^^^^^^^^^
Loading

0 comments on commit 7e0fc30

Please sign in to comment.