Add AttrItemMeta for handling attribute item attathced with once

Snowapril · Snowapril · commit 97f3434c0a4f · 2021-11-12T21:40:26.000+09:00
Signed-off-by: snowapril &lt;sinjihng@gmail.com&gt;
diff --git a/derive/src/pymodule.rs b/derive/src/pymodule.rs
@@ -1,7 +1,7 @@
 use crate::error::Diagnostic;
 use crate::util::{
-    iter_use_idents, path_eq, pyclass_ident_and_attrs, text_signature, AttributeExt, ClassItemMeta,
-    ContentItem, ContentItemInner, ErrorVec, ItemMeta, ItemNursery, SimpleItemMeta,
+    iter_use_idents, pyclass_ident_and_attrs, text_signature, AttrItemMeta, AttributeExt,
+    ClassItemMeta, ContentItem, ContentItemInner, ErrorVec, ItemMeta, ItemNursery, SimpleItemMeta,
     ALL_ALLOWED_NAMES,
 };
 use proc_macro2::TokenStream;
@@ -37,41 +37,6 @@ pub fn impl_pymodule(attr: AttributeArgs, module_item: Item) -> Result<TokenStre
     // collect to context
     for item in items.iter_mut() {
         let r = item.try_split_attr_mut(|attrs, item| {
-            // If attribute is #[pyattr] and item is ItemFn, then
-            // wrapping it with static_cell for preventing it from using it as function
-            if attrs.iter().any(|attr| {
-                path_eq(&attr.path, "pyattr")
-                    && if let Ok(syn::Meta::List(l)) = attr.parse_meta() {
-                        l.nested
-                            .into_iter()
-                            .any(|n| n.get_ident().map_or(false, |p| p == "once"))
-                    } else {
-                        false
-                    }
-            }) {
-                if let Item::Fn(syn::ItemFn { sig, block, .. }) = item {
-                    let stmts = &block.stmts;
-                    let return_type = match &sig.output {
-                        syn::ReturnType::Default => {
-                            unreachable!("#[pyattr] attached function must have return type.")
-                        }
-                        syn::ReturnType::Type(_, ty) => ty,
-                    };
-                    let stmt: syn::Stmt = parse_quote! {
-                        {
-                            rustpython_common::static_cell! {
-                                static ERROR: #return_type;
-                            }
-                            ERROR
-                                .get_or_init(|| {
-                                    #(#stmts)*
-                                })
-                                .clone()
-                        }
-                    };
-                    block.stmts = vec![stmt];
-                }
-            }
             let (pyitems, cfgs) = attrs_to_module_items(attrs, new_module_item)?;
             for pyitem in pyitems.iter().rev() {
                 let r = pyitem.gen_module_item(ModuleItemArgs {
@@ -276,7 +241,7 @@ impl ContentItem for AttributeItem {
 }
 
 struct ModuleItemArgs<'a> {
-    item: &'a Item,
+    item: &'a mut Item,
     attrs: &'a mut Vec<Attribute>,
     context: &'a mut ModuleContext,
     cfgs: &'a [Attribute],
@@ -405,15 +370,36 @@ impl ModuleItem for AttributeItem {
     fn gen_module_item(&self, args: ModuleItemArgs<'_>) -> Result<()> {
         let cfgs = args.cfgs.to_vec();
         let attr = args.attrs.remove(self.index());
-        let get_py_name = |attr: &Attribute, ident: &Ident| -> Result<_> {
-            let item_meta = SimpleItemMeta::from_attr(ident.clone(), attr)?;
-            let py_name = item_meta.simple_name()?;
-            Ok(py_name)
-        };
         let (py_name, tokens) = match args.item {
-            Item::Fn(syn::ItemFn { sig, .. }) => {
+            Item::Fn(syn::ItemFn { sig, block, .. }) => {
                 let ident = &sig.ident;
-                let py_name = get_py_name(&attr, ident)?;
+                // If `once` keyword is in #[pyattr],
+                // wrapping it with static_cell for preventing it from using it as function
+                let attr_meta = AttrItemMeta::from_attr(ident.clone(), &attr)?;
+                if attr_meta.inner()._bool("once")? {
+                    let stmts = &block.stmts;
+                    let return_type = match &sig.output {
+                        syn::ReturnType::Default => {
+                            unreachable!("#[pyattr] attached function must have return type.")
+                        }
+                        syn::ReturnType::Type(_, ty) => ty,
+                    };
+                    let stmt: syn::Stmt = parse_quote! {
+                        {
+                            rustpython_common::static_cell! {
+                                static ERROR: #return_type;
+                            }
+                            ERROR
+                                .get_or_init(|| {
+                                    #(#stmts)*
+                                })
+                                .clone()
+                        }
+                    };
+                    block.stmts = vec![stmt];
+                }
+
+                let py_name = attr_meta.simple_name()?;
                 (
                     py_name.clone(),
                     quote_spanned! { ident.span() =>
@@ -422,7 +408,8 @@ impl ModuleItem for AttributeItem {
                 )
             }
             Item::Const(syn::ItemConst { ident, .. }) => {
-                let py_name = get_py_name(&attr, ident)?;
+                let item_meta = SimpleItemMeta::from_attr(ident.clone(), &attr)?;
+                let py_name = item_meta.simple_name()?;
                 (
                     py_name.clone(),
                     quote_spanned! { ident.span() =>
diff --git a/derive/src/util.rs b/derive/src/util.rs
@@ -226,6 +226,19 @@ pub(crate) trait ItemMeta: Sized {
 pub(crate) struct SimpleItemMeta(pub ItemMetaInner);
 
 impl ItemMeta for SimpleItemMeta {
+    const ALLOWED_NAMES: &'static [&'static str] = &["name"];
+
+    fn from_inner(inner: ItemMetaInner) -> Self {
+        Self(inner)
+    }
+    fn inner(&self) -> &ItemMetaInner {
+        &self.0
+    }
+}
+
+pub(crate) struct AttrItemMeta(pub ItemMetaInner);
+
+impl ItemMeta for AttrItemMeta {
     const ALLOWED_NAMES: &'static [&'static str] = &["name", "once"];
 
     fn from_inner(inner: ItemMetaInner) -> Self {
