taitan_orm_parser/info_parser/
attr_parser.rs

1use quote::ToTokens;
2use std::borrow::Cow;
3use syn::punctuated::Punctuated;
4use syn::{Attribute, Expr, Lit, Meta, Path, Token};
5
6// is_attr  (&attr, name) -> bool
7// has_attr (&[attrs], name) -> bool
8// get_attr (&[attrs], name) -> Option<Attribute>
9// get_attrs(&[attrs], name) -> Vec<Attribute>
10
11// parse       (&attr) -> Option<NamedAttribute>
12// parse_one   (&attr) -> NamedAttribute
13// parse_list  (&attr) -> Vec<NamedAttribute>
14
15// extract           (&[attrs], name) -> Option<NamedAttribute>
16// extract_one       (&[attrs], name) -> NamedAttribute
17// extract_list      (&[attrs], name) -> Vec<NamedAttribute>
18// extract_multi_one (&[attrs], name) -> Vec<NamedAttribute>
19// extract_multi_list(&[attrs], name) -> Vec<Vec<NamedAttribute>>
20
21pub struct AttrParser;
22
23#[derive(Debug, Clone, PartialOrd, PartialEq)]
24pub struct NamedAttribute<'a> {
25    pub name: Cow<'a, str>,
26    pub values: Vec<Cow<'a, str>>,
27}
28#[derive(Debug, Clone, PartialOrd, PartialEq)]
29pub struct NamedAttributes<'a> {
30    pub name: Cow<'a, str>,
31    pub attrs: Vec<NamedAttribute<'a>>,
32}
33
34impl<'a> NamedAttribute<'a> {
35    pub fn from_str<N, F>(name: N, val_str: F) -> Self
36    where
37        N: Into<Cow<'a, str>>,
38        F: Into<Cow<'a, str>>,
39    {
40        let value_str: String = val_str.into().into_owned();
41        let values: Vec<Cow<'a, str>> = value_str
42            .split(|c| c == ' ' || c == ',')
43            .filter(|s| !s.is_empty())
44            .map(|s| Cow::Owned(s.to_string()))
45            .collect();
46        Self {
47            name: name.into(),
48            values,
49        }
50    }
51
52    pub fn from_origin<N, F>(name: N, val_str: F) -> Self
53    where
54        N: Into<Cow<'a, str>>,
55        F: Into<Cow<'a, str>>,
56    {
57        Self {
58            name: name.into(),
59            values: vec![val_str.into()],
60        }
61    }
62    pub fn from_vec<N>(name: N, values: Vec<Cow<'a, str>>) -> Self
63    where
64        N: Into<Cow<'a, str>>,
65    {
66        Self {
67            name: name.into(),
68            values,
69        }
70    }
71
72    pub fn get_single_value(&self) -> &Cow<'a, str> {
73        if self.values.len() != 1 {
74            panic!("attribute must have a single value");
75        }
76        &self.values[0]
77    }
78
79}
80impl AttrParser {
81    pub fn is_attr(attr: &Attribute, name: &str) -> bool {
82        let path: &Path = attr.path();
83        let path_ident = path.get_ident().unwrap();
84        let attr_path_name = path_ident.to_string();
85        attr_path_name == name
86    }
87    pub fn has_attr(attrs: &[Attribute], name: &str) -> bool {
88        attrs.iter().any(|attr| Self::is_attr(attr, name))
89    }
90
91    pub fn get_attr<'a>(attrs: &'a [Attribute], name: &'a str) -> Option<&'a Attribute> {
92        attrs.iter().find(|a| Self::is_attr(a, name))
93    }
94    pub fn get_attrs<'a>(attrs: &'a [Attribute], name: &'a str) -> Vec<&'a Attribute> {
95        attrs
96            .iter()
97            .filter(|a| Self::is_attr(a, name))
98            .collect()
99    }
100
101
102    pub fn parse<'a>(attr: &'a Attribute) -> Option<NamedAttribute<'a>> {
103        let path: &Path = attr.path();
104        let attr_name = path.get_ident().unwrap().to_string();
105        match &attr.meta {
106            Meta::NameValue(name_value) => match &name_value.value {
107                Expr::Lit(s) => match &s.lit {
108                    Lit::Str(s) => Some(NamedAttribute::from_str(attr_name, s.value())),
109                    _ => None,
110                },
111                Expr::Path(expr_path) => {
112                    let segments = &expr_path.path.segments;
113                    if segments.is_empty() {
114                        return None;
115                    }
116                    let first_segment = segments.first().unwrap();
117                    Some(NamedAttribute::from_str(
118                        attr_name,
119                        first_segment.ident.to_string(),
120                    ))
121                }
122                _ => None,
123            },
124            Meta::List(list) => {
125                if let Ok(expr) = list.parse_args::<Expr>() {
126                    return match expr {
127                        Expr::Lit(expr_lit) => {
128                            // 处理 #[attr("val")]
129                            if let Lit::Str(s) = expr_lit.lit {
130                                return Some(NamedAttribute::from_str(attr_name, s.value()));
131                            }
132                            None
133                        }
134                        Expr::Path(expr_path) => {
135                            let segments = expr_path.path.segments;
136                            if segments.is_empty() {
137                                return None;
138                            }
139                            let first_segment = segments.first().unwrap();
140                            Some(NamedAttribute::from_str(
141                                attr_name,
142                                first_segment.ident.to_string(),
143                            ))
144                        }
145                        _ => None,
146                    };
147                }
148                if let Ok(value_list) =
149                    list.parse_args_with(Punctuated::<Expr, Token![,]>::parse_terminated)
150                {
151                    let values = value_list
152                        .into_iter()
153                        .map(|v| Cow::Owned(v.to_token_stream().to_string()))
154                        .collect();
155                    return Some(NamedAttribute::from_vec(attr_name, values));
156                }
157
158                None
159            }
160            _ => None,
161        }
162    }
163
164    pub fn parse_one<'a>(attrs: &'a Attribute) -> NamedAttribute<'a> {
165        let attr_opt = Self::parse(attrs);
166        if let Some(attr) = attr_opt {
167            attr
168        } else {
169            panic!("cannot parse attribute")
170        }
171    }
172
173    pub fn parse_one_single<'a>(attrs: &'a Attribute) -> NamedAttribute<'a> {
174        let attr_opt = Self::parse(attrs);
175        if let Some(attr) = attr_opt {
176            if attr.values.len() != 1 {
177                panic!("cannot parse attribute to one single value")
178            }
179            attr
180        } else {
181            panic!("cannot parse attribute")
182        }
183    }
184
185    pub fn parse_list<'a>(attr: &'a Attribute) -> Vec<NamedAttribute<'a>> {
186        let path: &Path = attr.path();
187        let attr_name = path.get_ident().unwrap().to_string();
188        match &attr.meta {
189            Meta::NameValue(name_value) => match &name_value.value {
190                Expr::Lit(s) => match &s.lit {
191                    Lit::Str(s) => vec![NamedAttribute::from_str(attr_name, s.value())],
192                    _ => vec![],
193                },
194                Expr::Path(expr_path) => {
195                    let segments = &expr_path.path.segments;
196                    if segments.is_empty() {
197                        return vec![];
198                    }
199                    let first_segment = segments.first().unwrap();
200                    vec![NamedAttribute::from_str(
201                        attr_name,
202                        first_segment.ident.to_string(),
203                    )]
204                }
205                _ => vec![],
206            },
207            Meta::List(list) => {
208                if let Ok(expr_list) =
209                    list.parse_args_with(Punctuated::<Expr, Token![,]>::parse_terminated)
210                {
211                    let values: Vec<NamedAttribute> = expr_list
212                        .into_iter()
213                        .map(|expr| {
214                            if let Expr::Assign(assign) = expr {
215                                let name = assign.left.into_token_stream().to_string();
216                                let values_str = assign.right.into_token_stream().to_string();
217                                // TODO: this is hard coding, need to fix
218                                return if name == "generated" {
219                                    let values_inner_str = extract_quote_string(values_str.as_str());
220                                    Some(NamedAttribute::from_origin(name, values_inner_str))
221                                } else {
222                                    let values_inner_str = extract_inner_string(values_str.as_str());
223                                    Some(NamedAttribute::from_str(name, values_inner_str))
224                                }
225                            }
226                            if let Expr::Call(call_expr) = expr {
227                                let name = call_expr.func.to_token_stream().to_string();
228                                let args = call_expr.args.into_token_stream().to_string();
229                                return Some(NamedAttribute::from_str(name, args));
230                            }
231                            return None;
232                        })
233                        .filter(|v| v.is_some())
234                        .map(|v| v.unwrap())
235                        .collect();
236                    return values;
237                }
238
239
240                if let Ok(expr) = list.parse_args::<Expr>() {
241                    return match expr {
242                        Expr::Lit(expr_lit) => {
243                            // 处理 #[attr("val")]
244                            if let Lit::Str(s) = expr_lit.lit {
245                                return vec![NamedAttribute::from_str(attr_name, s.value())];
246                            }
247                            vec![]
248                        }
249                        Expr::Path(expr_path) => {
250                            let segments = expr_path.path.segments;
251                            if segments.is_empty() {
252                                return vec![];
253                            }
254                            let first_segment = segments.first().unwrap();
255                            return vec![NamedAttribute::from_str(
256                                attr_name,
257                                first_segment.ident.to_string(),
258                            )];
259                        }
260                        _ => vec![],
261                    };
262                }
263                if let Ok(value_list) =
264                    list.parse_args_with(Punctuated::<Expr, Token![,]>::parse_terminated)
265                {
266                    let values = value_list
267                        .into_iter()
268                        .map(|v| Cow::Owned(v.to_token_stream().to_string()))
269                        .collect();
270                    return vec![NamedAttribute::from_vec(attr_name, values)];
271                }
272                vec![]
273            }
274            _ => vec![],
275        }
276    }
277
278    pub fn extract<'a>(attrs: &'a [Attribute], name: &'a str) -> Option<NamedAttribute<'a>> {
279        let attr_opt = attrs.iter().find(|a| Self::is_attr(a, name));
280        if let Some(attr) = attr_opt {
281            let attr_opt = Self::parse(&attr);
282            if let Some(named) = attr_opt {
283                Some(named)
284            } else {
285                panic!("cannot parse attribute")
286            }
287        } else {
288            None
289        }
290    }
291
292    pub fn extract_one<'a>(attrs: &'a [Attribute], name: &'a str) -> NamedAttribute<'a> {
293        let attr_opt = Self::extract(attrs, name);
294        if let Some(attr) = attr_opt {
295            attr
296        } else {
297            panic!("cannot extract attribute")
298        }
299    }
300
301    pub fn extract_one_single<'a>(attrs: &'a [Attribute], name: &'a str) -> NamedAttribute<'a> {
302        let attr_opt = Self::extract(attrs, name);
303        if let Some(attr) = attr_opt {
304            if attr.values.len() != 1 {
305                panic!("cannot extract attribute to one single value")
306            }
307            attr
308        } else {
309            panic!("cannot extract attribute")
310        }
311    }
312
313    pub fn extract_list<'a>(attrs: &'a [Attribute], name: &str) -> Vec<NamedAttribute<'a>> {
314        let attr_opt = attrs.iter().find(|a| Self::is_attr(a, name));
315        if let Some(attr) = attr_opt {
316            Self::parse_list(&attr)
317        } else {
318            Vec::new()
319        }
320    }
321
322    pub fn extract_multi_one<'a>(attrs: &'a [Attribute], name: &'a str) -> Vec<NamedAttribute<'a>> {
323        let attrs: Vec<&Attribute> = attrs
324            .iter()
325            .filter(|a| Self::is_attr(a, name))
326            .collect();
327        let mut attr_list = Vec::new();
328        for attr in attrs {
329            if let Some(named) = Self::parse(&attr) {
330                attr_list.push(named.clone());
331            }
332        }
333        attr_list
334    }
335
336    pub fn extract_multi_list<'a>(attrs: &'a [Attribute], name: &'a str) -> Vec<NamedAttributes<'a>> {
337        let attrs: Vec<&Attribute> = attrs
338            .iter()
339            .filter(|a| Self::is_attr(a, name))
340            .collect();
341        let mut attr_list = Vec::new();
342        for attr in attrs {
343            let list = Self::parse_list(&attr);
344            let named_attrs = NamedAttributes {
345                name: Cow::Borrowed(name),
346                attrs: list
347            };
348            attr_list.push(named_attrs);
349        }
350        attr_list
351    }
352}
353
354fn extract_quote_string(s: &str) -> String {
355    let trimmed = s.trim();
356    if trimmed.starts_with('"') && trimmed.ends_with('"') {
357        trimmed[1..trimmed.len() - 1].trim().to_string()
358    } else {
359        trimmed.to_string()
360    }
361}
362fn extract_inner_string(s: &str) -> String {
363    let trimmed = s.trim();
364    if trimmed.starts_with('(') && trimmed.ends_with(')') {
365        trimmed[1..trimmed.len() - 1].trim().to_string()
366    } else if trimmed.starts_with('"') && trimmed.ends_with('"') {
367        trimmed[1..trimmed.len() - 1].trim().to_string()
368    } else {
369        trimmed.to_string()
370    }
371}
taitan_orm_parser/info_parser/attr_parser.rs

taitan_orm_parser/info_parser/
attr_parser.rs
