virtue_next/parse/
attributes.rs

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