chore(craft-macros): Improve the robustness of extracting attributes.

salvo-rs · Sep 29, 2024 · cdeea13 · cdeea13
1 parent 072cded
commit cdeea13
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 18 deletions.
diff --git a/crates/craft-macros/Cargo.toml b/crates/craft-macros/Cargo.toml
@@ -25,10 +25,11 @@ proc-macro-crate.workspace = true
 proc-macro2.workspace = true
 quote.workspace = true
 syn = { workspace = true, features = ["full", "parsing"] }
+regex.workspace = true
 
 [dev-dependencies]
 salvo = { path = "../salvo", features = ["oapi"] }
 tokio.workspace = true
 
 [lints]
-workspace = true
+workspace = true
diff --git a/crates/craft-macros/examples/example-add.rs b/crates/craft-macros/examples/example-add.rs
@@ -33,7 +33,7 @@ impl Opts {
     }
     /// doc line 5
     /// doc line 6
-    #[craft(endpoint(responses((status_code = 400, description = "Wrong request parameters."))))]
+    #[craft(endpoint(responses((status_code = 400, description = "[(Wrong)] request parameters."))))]
     pub fn add3(left: QueryParam<i64>, right: QueryParam<i64>) -> String {
         (*left + *right).to_string()
     }

diff --git a/crates/craft-macros/src/craft.rs b/crates/craft-macros/src/craft.rs
@@ -1,6 +1,7 @@
 use crate::utils::salvo_crate;
 use proc_macro2::{Span, TokenStream};
 use quote::{quote, ToTokens};
+use regex::Regex;
 use syn::parse::Parser;
 use syn::{
     parse_quote, Attribute, FnArg, Generics, Ident, ImplItem, ImplItemFn, Item, Token, Type,
@@ -28,6 +29,9 @@ pub(crate) fn generate(input: Item) -> syn::Result<TokenStream> {
     }
 }
 
+const REGEX_STR: &'static str =
+    r#"(?s)#\s*\[\s*craft\s*\(\s*(?P<name>handler|endpoint)\s*(?P<content>\(.*\))?\s*\)\s*\]"#;
+
 fn take_method_macro(item_fn: &mut ImplItemFn) -> syn::Result<Option<Attribute>> {
     let mut index: Option<usize> = None;
     let mut new_attr: Option<Attribute> = None;
@@ -38,26 +42,31 @@ fn take_method_macro(item_fn: &mut ImplItemFn) -> syn::Result<Option<Attribute>>
         }) {
             continue;
         }
-        if let Some((_, last)) = attr.to_token_stream().to_string().split_once("craft(") {
-            if let Some(last) = last.strip_suffix(")]") {
-                let ts: Option<TokenStream> = if last == "handler" || last.starts_with("handler(") {
-                    Some(format!("#[{}::{last}]", salvo_crate()).parse()?)
-                } else if last == "endpoint" || last.starts_with("endpoint(") {
-                    Some(format!("#[{}::oapi::{last}]", salvo_crate()).parse()?)
-                } else {
-                    None
+        let re = Regex::new(REGEX_STR).unwrap();
+        let attr_str = attr.to_token_stream().to_string().trim().to_owned();
+        if let Some(caps) = re.captures(&attr_str) {
+            if let Some(name) = caps.name("name") {
+                let name = name.as_str();
+                let content = caps
+                    .name("content")
+                    .map(|c| c.as_str().to_string())
+                    .unwrap_or_default();
+                let ts: TokenStream = match name {
+                    "handler" => format!("#[{}::{name}{content}]", salvo_crate()).parse()?,
+                    "endpoint" => format!("#[{}::oapi::{name}{content}]", salvo_crate()).parse()?,
+                    _ => {
+                        unreachable!()
+                    }
                 };
-                if let Some(ts) = ts {
-                    new_attr = Attribute::parse_outer.parse2(ts)?.into_iter().next();
-                    index = Some(idx);
-                    continue;
-                }
+                new_attr = Attribute::parse_outer.parse2(ts)?.into_iter().next();
+                index = Some(idx);
+                continue;
             }
         }
         return Err(syn::Error::new_spanned(
-            item_fn,
-            "The attribute macro #[craft] on a method must be filled with sub-attributes, such as '#[craft(handler)]', '#[craft(endpoint)]', or '#[craft(endpoint(...))]'.",
-        ));
+                item_fn,
+                format!("The attribute macro `{attr_str}` on a method must be filled with sub-attributes, such as '#[craft(handler)]', '#[craft(endpoint)]', or '#[craft(endpoint(...))]'."),
+            ));
     }
     if let Some(index) = index {
         item_fn.attrs.remove(index);