Skip to main content

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 mut ts1 = token_stream("struct Foo;");
81    let stream = &mut ts1;
82    assert!(
83        Attribute::try_take(AttributeLocation::Container, stream)
84            .unwrap()
85            .is_empty()
86    );
87    match stream.next().unwrap() {
88        TokenTree::Ident(i) => assert_eq!(i, "struct"),
89        x => panic!("Expected ident, found {:?}", x),
90    }
91
92    let mut ts2 = token_stream("#[cfg(test)] struct Foo;");
93    let stream = &mut ts2;
94    assert!(
95        !Attribute::try_take(AttributeLocation::Container, stream)
96            .unwrap()
97            .is_empty()
98    );
99    match stream.next().unwrap() {
100        TokenTree::Ident(i) => assert_eq!(i, "struct"),
101        x => panic!("Expected ident, found {:?}", x),
102    }
103}
104
105/// Helper trait for [`AttributeAccess`] methods.
106///
107/// This can be implemented on your own type to make parsing easier.
108///
109/// Some functions that can make your life easier:
110/// - [`utils::parse_tagged_attribute`] is a helper for parsing attributes in the format of `#[prefix(...)]`
111///
112/// [`AttributeAccess`]: trait.AttributeAccess.html
113/// [`utils::parse_tagged_attribute`]: ../utils/fn.parse_tagged_attribute.html
114pub trait FromAttribute: Sized {
115    /// 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.
116    fn parse(group: &Group) -> Result<Option<Self>>;
117}
118
119/// Bring useful methods to access attributes of an element.
120pub trait AttributeAccess {
121    /// Check to see if has the given attribute. See [`FromAttribute`] for more information.
122    ///
123    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
124    fn has_attribute<T: FromAttribute + PartialEq<T>>(&self, attrib: T) -> Result<bool>;
125
126    /// Returns the first attribute that returns `Some(Self)`. See [`FromAttribute`] for more information.
127    ///
128    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
129    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>>;
130}
131
132impl AttributeAccess for Vec<Attribute> {
133    fn has_attribute<T: FromAttribute + PartialEq<T>>(&self, attrib: T) -> Result<bool> {
134        for attribute in self.iter() {
135            let parsed = T::parse(&attribute.tokens)?;
136            if let Some(attribute) = parsed {
137                if attribute == attrib {
138                    return Ok(true);
139                }
140            }
141        }
142        Ok(false)
143    }
144
145    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>> {
146        for attribute in self.iter() {
147            if let Some(attribute) = T::parse(&attribute.tokens)? {
148                return Ok(Some(attribute));
149            }
150        }
151        Ok(None)
152    }
153}