Skip to main content

virtue_next/parse/
attributes.rs

1use super::utils::assume_group;
2use super::utils::consume_punct_if;
3use crate::Error;
4use crate::Result;
5use crate::prelude::Delimiter;
6use crate::prelude::Group;
7use crate::prelude::Punct;
8use crate::prelude::TokenTree;
9use std::iter::Peekable;
10
11/// An attribute for the given struct, enum, field, etc
12#[derive(Debug, Clone)]
13#[non_exhaustive]
14pub struct Attribute {
15    /// The location this attribute was parsed at
16    pub location: AttributeLocation,
17    /// The punct token of the attribute. This will always be `Punct('#')`
18    pub punct: Punct,
19    /// The group of tokens of the attribute. You can parse this to get your custom attributes.
20    pub tokens: Group,
21}
22
23/// The location an attribute can be found at
24#[derive(PartialEq, Eq, Debug, Hash, Copy, Clone)]
25#[non_exhaustive]
26pub enum AttributeLocation {
27    /// The attribute is on a container, which will be either a `struct` or an `enum`
28    Container,
29    /// The attribute is on an enum variant
30    Variant,
31    /// The attribute is on a field, which can either be a struct field or an enum variant field
32    /// ```ignore
33    /// struct Foo {
34    ///     #[attr] // here
35    ///     pub a: u8
36    /// }
37    /// struct Bar {
38    ///     Baz {
39    ///         #[attr] // or here
40    ///         a: u8
41    ///     }
42    /// }
43    /// ```
44    Field,
45}
46
47impl Attribute {
48    pub(crate) fn try_take(
49        location: AttributeLocation,
50        input: &mut Peekable<impl Iterator<Item = TokenTree>>,
51    ) -> Result<Vec<Self>> {
52        let mut result = Vec::new();
53
54        while let Some(punct) = consume_punct_if(input, '#') {
55            match input.peek() {
56                | Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
57                    let group = assume_group(input.next());
58                    result.push(Self {
59                        location,
60                        punct,
61                        tokens: group,
62                    });
63                },
64                | Some(TokenTree::Group(g)) => {
65                    return Err(Error::InvalidRustSyntax {
66                        span: g.span(),
67                        expected: format!("[] bracket, got {:?}", g.delimiter()),
68                    });
69                },
70                | Some(TokenTree::Punct(p)) if p.as_char() == '#' => {
71                    // sometimes with empty lines of doc comments, we get two #'s in a row
72                    // Just ignore this
73                },
74                | token => return Error::wrong_token(token, "[] group or next # attribute"),
75            }
76        }
77        Ok(result)
78    }
79}
80
81#[test]
82fn test_attributes_try_take() {
83    use crate::token_stream;
84
85    let mut ts1 = token_stream("struct Foo;");
86    let stream = &mut ts1;
87    assert!(
88        Attribute::try_take(AttributeLocation::Container, stream)
89            .unwrap()
90            .is_empty()
91    );
92    match stream.next().unwrap() {
93        | TokenTree::Ident(i) => assert_eq!(i, "struct"),
94        | x => panic!("Expected ident, found {:?}", x),
95    }
96
97    let mut ts2 = token_stream("#[cfg(test)] struct Foo;");
98    let stream = &mut ts2;
99    assert!(
100        !Attribute::try_take(AttributeLocation::Container, stream)
101            .unwrap()
102            .is_empty()
103    );
104    match stream.next().unwrap() {
105        | TokenTree::Ident(i) => assert_eq!(i, "struct"),
106        | x => panic!("Expected ident, found {:?}", x),
107    }
108}
109
110/// Helper trait for [`AttributeAccess`] methods.
111///
112/// This can be implemented on your own type to make parsing easier.
113///
114/// Some functions that can make your life easier:
115/// - [`utils::parse_tagged_attribute`] is a helper for parsing attributes in the format of `#[prefix(...)]`
116///
117/// [`AttributeAccess`]: trait.AttributeAccess.html
118/// [`utils::parse_tagged_attribute`]: ../utils/fn.parse_tagged_attribute.html
119pub trait FromAttribute: Sized {
120    /// Try to parse the given group into your own type. Return `Ok(None)` if the parsing failed or if the attribute was not this type.
121    ///
122    /// # Errors
123    ///
124    /// Returns an error if the operation fails.
125    fn parse(group: &Group) -> Result<Option<Self>>;
126}
127
128/// Bring useful methods to access attributes of an element.
129pub trait AttributeAccess {
130    /// Check to see if has the given attribute. See [`FromAttribute`] for more information.
131    ///
132    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error if the operation fails.
137    fn has_attribute<T: FromAttribute + PartialEq<T>>(
138        &self,
139        attrib: T,
140    ) -> Result<bool>;
141
142    /// Returns the first attribute that returns `Some(Self)`. See [`FromAttribute`] for more information.
143    ///
144    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
145    ///
146    /// # Errors
147    ///
148    /// Returns an error if the operation fails.
149    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>>;
150}
151
152impl AttributeAccess for Vec<Attribute> {
153    fn has_attribute<T: FromAttribute + PartialEq<T>>(
154        &self,
155        attrib: T,
156    ) -> Result<bool> {
157        for attribute in self {
158            let parsed = T::parse(&attribute.tokens)?;
159            if let Some(attribute) = parsed
160                && attribute == attrib
161            {
162                return Ok(true);
163            }
164        }
165        Ok(false)
166    }
167
168    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>> {
169        for attribute in self {
170            if let Some(attribute) = T::parse(&attribute.tokens)? {
171                return Ok(Some(attribute));
172            }
173        }
174        Ok(None)
175    }
176}