Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow inline struct, enum, fn and mod inside of declarative modules #3902

Merged
merged 2 commits into from
Mar 6, 2024
Merged

Allow inline struct, enum, fn and mod inside of declarative modules #3902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6b30a07
Inline struct, enum, fn and mod inside of declarative modules
Tpt Feb 23, 2024
bc8f777
remove news fragment
davidhewitt Mar 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
161 changes: 135 additions & 26 deletions pyo3-macros-backend/src/module.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,19 +114,10 @@ pub fn pymodule_module_impl(mut module: syn::ItemMod) -> Result<TokenStream> {
for item in &mut *items {
match item {
Item::Use(item_use) => {
let mut is_pyo3 = false;
item_use.attrs.retain(|attr| {
let found = attr.path().is_ident("pymodule_export");
is_pyo3 |= found;
!found
});
if is_pyo3 {
let cfg_attrs = item_use
.attrs
.iter()
.filter(|attr| attr.path().is_ident("cfg"))
.cloned()
.collect::<Vec<_>>();
let is_pymodule_export =
find_and_remove_attribute(&mut item_use.attrs, "pymodule_export");
if is_pymodule_export {
let cfg_attrs = get_cfg_attributes(&item_use.attrs);
extract_use_items(
&item_use.tree,
&cfg_attrs,
Expand All @@ -136,23 +127,116 @@ pub fn pymodule_module_impl(mut module: syn::ItemMod) -> Result<TokenStream> {
}
}
Item::Fn(item_fn) => {
let mut is_module_init = false;
item_fn.attrs.retain(|attr| {
let found = attr.path().is_ident("pymodule_init");
is_module_init |= found;
!found
});
if is_module_init {
ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one pymodule_init may be specified");
let ident = &item_fn.sig.ident;
ensure_spanned!(
!has_attribute(&item_fn.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
let is_pymodule_init =
find_and_remove_attribute(&mut item_fn.attrs, "pymodule_init");
let ident = &item_fn.sig.ident;
if is_pymodule_init {
ensure_spanned!(
!has_attribute(&item_fn.attrs, "pyfunction"),
item_fn.span() => "`#[pyfunction]` cannot be used alongside `#[pymodule_init]`"
);
ensure_spanned!(pymodule_init.is_none(), item_fn.span() => "only one `#[pymodule_init]` may be specified");
pymodule_init = Some(quote! { #ident(module)?; });
} else {
bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
} else if has_attribute(&item_fn.attrs, "pyfunction") {
module_items.push(ident.clone());
module_items_cfg_attrs.push(get_cfg_attributes(&item_fn.attrs));
}
}
item => {
bail_spanned!(item.span() => "only 'use' statements and and pymodule_init functions are allowed in #[pymodule]")
Item::Struct(item_struct) => {
ensure_spanned!(
!has_attribute(&item_struct.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
if has_attribute(&item_struct.attrs, "pyclass") {
module_items.push(item_struct.ident.clone());
module_items_cfg_attrs.push(get_cfg_attributes(&item_struct.attrs));
}
}
Item::Enum(item_enum) => {
ensure_spanned!(
!has_attribute(&item_enum.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
if has_attribute(&item_enum.attrs, "pyclass") {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice way to reuse the existing machinery!

module_items.push(item_enum.ident.clone());
module_items_cfg_attrs.push(get_cfg_attributes(&item_enum.attrs));
}
}
Item::Mod(item_mod) => {
ensure_spanned!(
!has_attribute(&item_mod.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
if has_attribute(&item_mod.attrs, "pymodule") {
module_items.push(item_mod.ident.clone());
module_items_cfg_attrs.push(get_cfg_attributes(&item_mod.attrs));
}
}
Item::ForeignMod(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Trait(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Const(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Static(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Macro(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::ExternCrate(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Impl(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::TraitAlias(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Type(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
Item::Union(item) => {
ensure_spanned!(
!has_attribute(&item.attrs, "pymodule_export"),
item.span() => "`#[pymodule_export]` may only be used on `use` statements"
);
}
_ => (),
}
}

Expand Down Expand Up @@ -352,6 +436,31 @@ fn get_pyfn_attr(attrs: &mut Vec<syn::Attribute>) -> syn::Result<Option<PyFnArgs
Ok(pyfn_args)
}

fn get_cfg_attributes(attrs: &[syn::Attribute]) -> Vec<syn::Attribute> {
attrs
.iter()
.filter(|attr| attr.path().is_ident("cfg"))
.cloned()
.collect()
}

fn find_and_remove_attribute(attrs: &mut Vec<syn::Attribute>, ident: &str) -> bool {
let mut found = false;
attrs.retain(|attr| {
if attr.path().is_ident(ident) {
found = true;
false
} else {
true
}
});
found
}

fn has_attribute(attrs: &[syn::Attribute], ident: &str) -> bool {
attrs.iter().any(|attr| attr.path().is_ident(ident))
}

enum PyModulePyO3Option {
Crate(CrateAttribute),
Name(NameAttribute),
Expand Down
36 changes: 32 additions & 4 deletions tests/test_declarative_module.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,8 @@ struct ValueClass {
#[pymethods]
impl ValueClass {
#[new]
fn new(value: usize) -> ValueClass {
ValueClass { value }
fn new(value: usize) -> Self {
Self { value }
}
}

Expand Down Expand Up @@ -48,6 +48,33 @@ mod declarative_module {
#[pymodule_export]
use super::{declarative_module2, double, MyError, ValueClass as Value};

#[pymodule]
mod inner {
use super::*;

#[pyfunction]
fn triple(x: usize) -> usize {
x * 3
}

#[pyclass]
struct Struct;

#[pymethods]
impl Struct {
#[new]
fn new() -> Self {
Self
}
}

#[pyclass]
enum Enum {
A,
B,
}
}

#[pymodule_init]
fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add("double2", m.getattr("double")?)
Expand All @@ -65,7 +92,6 @@ mod declarative_submodule {
use super::{double, double_value};
}

/// A module written using declarative syntax.
#[pymodule]
#[pyo3(name = "declarative_module_renamed")]
mod declarative_module2 {
Expand All @@ -84,7 +110,7 @@ fn test_declarative_module() {
);

py_assert!(py, m, "m.double(2) == 4");
py_assert!(py, m, "m.double2(3) == 6");
py_assert!(py, m, "m.inner.triple(3) == 9");
py_assert!(py, m, "m.declarative_submodule.double(4) == 8");
py_assert!(
py,
Expand All @@ -97,5 +123,7 @@ fn test_declarative_module() {
py_assert!(py, m, "not hasattr(m, 'LocatedClass')");
#[cfg(not(Py_LIMITED_API))]
py_assert!(py, m, "hasattr(m, 'LocatedClass')");
py_assert!(py, m, "isinstance(m.inner.Struct(), m.inner.Struct)");
py_assert!(py, m, "isinstance(m.inner.Enum.A, m.inner.Enum)");
})
}
2 changes: 1 addition & 1 deletion tests/ui/invalid_pymodule_trait.stderr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
error: only 'use' statements and and pymodule_init functions are allowed in #[pymodule]
error: `#[pymodule_export]` may only be used on `use` statements
--> tests/ui/invalid_pymodule_trait.rs:5:5
|
5 | #[pymodule_export]
Expand Down
2 changes: 1 addition & 1 deletion tests/ui/invalid_pymodule_two_pymodule_init.stderr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
error: only one pymodule_init may be specified
error: only one `#[pymodule_init]` may be specified
--> tests/ui/invalid_pymodule_two_pymodule_init.rs:11:5
|
11 | fn init2(m: &PyModule) -> PyResult<()> {
Expand Down