diff --git a/derive/src/pymodule.rs b/derive/src/pymodule.rs index aa39c2b15..dac487196 100644 --- a/derive/src/pymodule.rs +++ b/derive/src/pymodule.rs @@ -1,11 +1,11 @@ use crate::error::Diagnostic; use crate::util::{ - pyclass_ident_and_attrs, AttributeExt, ClassItemMeta, ContentItem, ContentItemInner, ErrorVec, - ItemMeta, ItemNursery, SimpleItemMeta, ALL_ALLOWED_NAMES, + iter_use_idents, pyclass_ident_and_attrs, AttributeExt, ClassItemMeta, ContentItem, + ContentItemInner, ErrorVec, ItemMeta, ItemNursery, SimpleItemMeta, ALL_ALLOWED_NAMES, }; use proc_macro2::TokenStream; use quote::{quote, quote_spanned, ToTokens}; -use syn::{parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Result, UseTree}; +use syn::{parse_quote, spanned::Spanned, Attribute, AttributeArgs, Ident, Item, Result}; use syn_ext::ext::*; #[derive(Default)] @@ -367,60 +367,28 @@ impl ModuleItem for AttributeItem { }, ) } - Item::Use(syn::ItemUse { - tree: UseTree::Path(path), - .. - }) => { - let ident = match &*path.tree { - UseTree::Name(name) => &name.ident, - UseTree::Rename(rename) => &rename.rename, - UseTree::Group(syn::UseGroup { items, .. }) => { - for item in items { - let ident = match item { - UseTree::Name(name) => &name.ident, - UseTree::Rename(rename) => &rename.rename, - other => { - return Err(self.new_syn_error( - other.span(), - "can only contains a simple use or a group of simple uses", - )); - } - }; - let item_meta = SimpleItemMeta::from_attr(ident.clone(), &attr)?; - if item_meta.optional_name().is_some() { - // this check actually doesn't need to be placed in loop - return Err(self.new_syn_error( - ident.span(), - "`name` attribute is not allowed for multiple use items", - )); - } - let py_name = ident.to_string(); - let tokens = quote! { - vm.__module_set_attr(&module, #py_name, vm.new_pyobj(#ident)).unwrap(); - }; - args.context.module_extend_items.add_item( - py_name, - cfgs.clone(), - tokens, - )?; - } - return Ok(()); - } - other => { + Item::Use(item) => { + return iter_use_idents(item, |ident, is_unique| { + let item_meta = SimpleItemMeta::from_attr(ident.clone(), &attr)?; + let py_name = if is_unique { + item_meta.simple_name()? + } else if item_meta.optional_name().is_some() { + // this check actually doesn't need to be placed in loop return Err(self.new_syn_error( - other.span(), - "can only contains a simple use or a group of simple uses", + ident.span(), + "`name` attribute is not allowed for multiple use items", )); - } - }; - - let py_name = get_py_name(&attr, &ident)?; - ( - py_name.clone(), - quote! { + } else { + ident.to_string() + }; + let tokens = quote! { vm.__module_set_attr(&module, #py_name, vm.new_pyobj(#ident)).unwrap(); - }, - ) + }; + args.context + .module_extend_items + .add_item(py_name, cfgs.clone(), tokens)?; + Ok(()) + }); } other => { return Err( diff --git a/derive/src/util.rs b/derive/src/util.rs index cbd7ec5ae..766b96cbf 100644 --- a/derive/src/util.rs +++ b/derive/src/util.rs @@ -1,7 +1,7 @@ use proc_macro2::{Span, TokenStream}; use quote::{quote, ToTokens}; use std::collections::HashMap; -use syn::{spanned::Spanned, Attribute, Ident, Meta, MetaList, NestedMeta, Path, Result}; +use syn::{spanned::Spanned, Attribute, Ident, Meta, MetaList, NestedMeta, Path, Result, UseTree}; use syn_ext::ext::{AttributeExt as SynAttributeExt, *}; use syn_ext::types::PunctuatedNestedMeta; @@ -456,3 +456,41 @@ macro_rules! iter_chain { $(.chain(::std::iter::once($it)))* }; } + +pub(crate) fn iter_use_idents(item_use: &syn::ItemUse, mut f: F) -> Result<()> +where + F: FnMut(&syn::Ident, bool) -> Result<()>, +{ + match &item_use.tree { + UseTree::Name(name) => f(&name.ident, true)?, + UseTree::Rename(rename) => f(&rename.rename, true)?, + UseTree::Path(path) => match &*path.tree { + UseTree::Name(name) => f(&name.ident, true)?, + UseTree::Rename(rename) => f(&rename.rename, true)?, + UseTree::Path(_) => unreachable!(), + other => iter_use_tree_idents(other, &mut f)?, + }, + other => iter_use_tree_idents(other, &mut f)?, + } + Ok(()) +} + +fn iter_use_tree_idents(tree: &syn::UseTree, f: &mut F) -> Result<()> +where + F: FnMut(&syn::Ident, bool) -> Result<()>, +{ + match tree { + UseTree::Name(name) => f(&name.ident, false)?, + UseTree::Rename(rename) => f(&rename.rename, false)?, + UseTree::Path(path) => iter_use_tree_idents(&*path.tree, f)?, + UseTree::Group(syn::UseGroup { items, .. }) => { + for subtree in items { + iter_use_tree_idents(subtree, f)?; + } + } + UseTree::Glob(glob) => { + return Err(syn::Error::new_spanned(glob, "#[py*] doesn't allow '*'")) + } + } + Ok(()) +}