Skip to main content

virtue_next/
utils.rs

1//! Utility functions
2use crate::{Error, prelude::*};
3
4/// Parse a tagged attribute. This is very helpful for implementing [`FromAttribute`].
5///
6/// A tagged attribute is an attribute in the form of `#[prefix(result)]`. This function will return `Some(result)` if the `prefix` matches.
7///
8/// The contents of the result can be either:
9/// - `ParsedAttribute::Tagged(Ident)`, e.g. `#[serde(skip)]` will be `Tagged("skip")`
10/// - `ParsedAttribute::Property(Ident, lit)`, e.g. `#[bincode(crate = "foo")]` will be `Property("crate", "foo")`
11///
12/// # Examples
13/// ```
14/// # use virtue::prelude::*;
15/// # use std::str::FromStr;
16/// # fn parse_token_stream_group(input: &'static str) -> Group {
17/// #     let token_stream: TokenStream = proc_macro2::TokenStream::from_str(input).unwrap().into();
18/// #     let mut iter = token_stream.into_iter();
19/// #     let Some(TokenTree::Punct(_)) = iter.next() else { panic!() };
20/// #     let Some(TokenTree::Group(group)) = iter.next() else { panic!() };
21/// #     group
22/// # }
23/// use virtue::utils::{parse_tagged_attribute, ParsedAttribute};
24///
25/// // The attribute being parsed
26/// let group: Group = parse_token_stream_group("#[prefix(result, foo = \"bar\")]");
27///
28/// let attributes = parse_tagged_attribute(&group, "prefix").unwrap().unwrap();
29/// let mut iter = attributes.into_iter();
30///
31/// // The stream will contain the contents of the `prefix(...)`
32/// match iter.next() {
33///     Some(ParsedAttribute::Tag(i)) => {
34///         assert_eq!(i.to_string(), String::from("result"));
35///     },
36///     x => panic!("Unexpected attribute: {:?}", x)
37/// }
38///  match iter.next() {
39///     Some(ParsedAttribute::Property(key, val)) => {
40///         assert_eq!(key.to_string(), String::from("foo"));
41///         assert_eq!(val.to_string(), String::from("\"bar\""));
42///     },
43///     x => panic!("Unexpected attribute: {:?}", x)
44/// }
45///
46/// ```
47pub fn parse_tagged_attribute(group: &Group, prefix: &str) -> Result<Option<Vec<ParsedAttribute>>> {
48    let stream = &mut group.stream().into_iter();
49    if let Some(TokenTree::Ident(attribute_ident)) = stream.next() {
50        #[allow(clippy::cmp_owned)] // clippy is wrong
51        if attribute_ident.to_string() == prefix {
52            if let Some(TokenTree::Group(group)) = stream.next() {
53                let mut result = Vec::new();
54                let mut stream = group.stream().into_iter().peekable();
55                while let Some(token) = stream.next() {
56                    match (token, stream.peek()) {
57                        (TokenTree::Ident(key), Some(TokenTree::Punct(p)))
58                            if p.as_char() == ',' =>
59                        {
60                            result.push(ParsedAttribute::Tag(key));
61                            stream.next();
62                        }
63                        (TokenTree::Ident(key), None) => {
64                            result.push(ParsedAttribute::Tag(key));
65                            stream.next();
66                        }
67                        (TokenTree::Ident(key), Some(TokenTree::Punct(p)))
68                            if p.as_char() == '=' =>
69                        {
70                            stream.next();
71                            if let Some(TokenTree::Literal(lit)) = stream.next() {
72                                result.push(ParsedAttribute::Property(key, lit));
73
74                                match stream.next() {
75                                    Some(TokenTree::Punct(p)) if p.as_char() == ',' => {}
76                                    None => {}
77                                    x => {
78                                        return Err(Error::custom_at_opt_token("Expected `,`", x));
79                                    }
80                                }
81                            }
82                        }
83                        (x, _) => {
84                            return Err(Error::custom_at(
85                                "Expected `key` or `key = \"val\"`",
86                                x.span(),
87                            ));
88                        }
89                    }
90                }
91
92                return Ok(Some(result));
93            }
94        }
95    }
96    Ok(None)
97}
98
99#[derive(Clone, Debug)]
100#[non_exhaustive]
101/// A parsed attribute. See [`parse_tagged_attribute`] for more information.
102pub enum ParsedAttribute {
103    /// A tag, created by parsing `#[prefix(foo)]`
104    Tag(Ident),
105    /// A property, created by parsing `#[prefix(foo = "bar")]`
106    Property(Ident, Literal),
107}
108
109#[test]
110fn test_parse_tagged_attribute() {
111    let group: Group = match crate::token_stream("[prefix(result, foo = \"bar\", baz)]").next() {
112        Some(TokenTree::Group(group)) => group,
113        x => panic!("Unexpected token {:?}", x),
114    };
115
116    let attributes = parse_tagged_attribute(&group, "prefix").unwrap().unwrap();
117    let mut iter = attributes.into_iter();
118
119    // The stream will contain the contents of the `prefix(...)`
120    match iter.next() {
121        Some(ParsedAttribute::Tag(i)) => {
122            assert_eq!(i.to_string(), String::from("result"));
123        }
124        x => panic!("Unexpected attribute: {:?}", x),
125    }
126    match iter.next() {
127        Some(ParsedAttribute::Property(key, val)) => {
128            assert_eq!(key.to_string(), String::from("foo"));
129            assert_eq!(val.to_string(), String::from("\"bar\""));
130        }
131        x => panic!("Unexpected attribute: {:?}", x),
132    }
133    match iter.next() {
134        Some(ParsedAttribute::Tag(i)) => {
135            assert_eq!(i.to_string(), String::from("baz"));
136        }
137        x => panic!("Unexpected attribute: {:?}", x),
138    }
139}