1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use super::utils::*;
use crate::prelude::{Delimiter, Group, Punct, TokenTree};
use crate::{Error, Result};
use std::iter::Peekable;

/// An attribute for the given struct, enum, field, etc
#[derive(Debug)]
#[non_exhaustive]
pub struct Attribute {
    /// The location this attribute was parsed at
    pub location: AttributeLocation,
    /// The punct token of the attribute. This will always be `Punct('#')`
    pub punct: Punct,
    /// The group of tokens of the attribute. You can parse this to get your custom attributes.
    pub tokens: Group,
}

/// The location an attribute can be found at
#[derive(PartialEq, Eq, Debug, Hash, Copy, Clone)]
#[non_exhaustive]
pub enum AttributeLocation {
    /// The attribute is on a container, which will be either a `struct` or an `enum`
    Container,
    /// The attribute is on an enum variant
    Variant,
    /// The attribute is on a field, which can either be a struct field or an enum variant field
    /// ```ignore
    /// struct Foo {
    ///     #[attr] // here
    ///     pub a: u8
    /// }
    /// struct Bar {
    ///     Baz {
    ///         #[attr] // or here
    ///         a: u8
    ///     }
    /// }
    /// ```
    Field,
}

impl Attribute {
    pub(crate) fn try_take(
        location: AttributeLocation,
        input: &mut Peekable<impl Iterator<Item = TokenTree>>,
    ) -> Result<Vec<Self>> {
        let mut result = Vec::new();

        while let Some(punct) = consume_punct_if(input, '#') {
            match input.peek() {
                Some(TokenTree::Group(g)) if g.delimiter() == Delimiter::Bracket => {
                    let group = assume_group(input.next());
                    result.push(Attribute {
                        location,
                        punct,
                        tokens: group,
                    });
                }
                Some(TokenTree::Group(g)) => {
                    return Err(Error::InvalidRustSyntax {
                        span: g.span(),
                        expected: format!("[] bracket, got {:?}", g.delimiter()),
                    });
                }
                Some(TokenTree::Punct(p)) if p.as_char() == '#' => {
                    // sometimes with empty lines of doc comments, we get two #'s in a row
                    // Just ignore this
                }
                token => return Error::wrong_token(token, "[] group or next # attribute"),
            }
        }
        Ok(result)
    }
}

#[test]
fn test_attributes_try_take() {
    use crate::token_stream;

    let stream = &mut token_stream("struct Foo;");
    assert!(Attribute::try_take(AttributeLocation::Container, stream)
        .unwrap()
        .is_empty());
    match stream.next().unwrap() {
        TokenTree::Ident(i) => assert_eq!(i, "struct"),
        x => panic!("Expected ident, found {:?}", x),
    }

    let stream = &mut token_stream("#[cfg(test)] struct Foo;");
    assert!(!Attribute::try_take(AttributeLocation::Container, stream)
        .unwrap()
        .is_empty());
    match stream.next().unwrap() {
        TokenTree::Ident(i) => assert_eq!(i, "struct"),
        x => panic!("Expected ident, found {:?}", x),
    }
}

/// Helper trait for [`AttributeAccess`] methods.
///
/// This can be implemented on your own type to make parsing easier.
///
/// Some functions that can make your life easier:
/// - [`utils::parse_tagged_attribute`] is a helper for parsing attributes in the format of `#[prefix(...)]`
///
/// [`AttributeAccess`]: trait.AttributeAccess.html
/// [`utils::parse_tagged_attribute`]: ../utils/fn.parse_tagged_attribute.html
pub trait FromAttribute: Sized {
    /// 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.
    fn parse(group: &Group) -> Result<Option<Self>>;
}

/// Bring useful methods to access attributes of an element.
pub trait AttributeAccess {
    /// Check to see if has the given attribute. See [`FromAttribute`] for more information.
    ///
    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
    fn has_attribute<T: FromAttribute + PartialEq<T>>(&self, attrib: T) -> Result<bool>;

    /// Returns the first attribute that returns `Some(Self)`. See [`FromAttribute`] for more information.
    ///
    /// **note**: Will immediately return `Err(_)` on the first error `T` returns.
    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>>;
}

impl AttributeAccess for Vec<Attribute> {
    fn has_attribute<T: FromAttribute + PartialEq<T>>(&self, attrib: T) -> Result<bool> {
        for attribute in self.iter() {
            if let Some(attribute) = T::parse(&attribute.tokens)? {
                if attribute == attrib {
                    return Ok(true);
                }
            }
        }
        Ok(false)
    }

    fn get_attribute<T: FromAttribute>(&self) -> Result<Option<T>> {
        for attribute in self.iter() {
            if let Some(attribute) = T::parse(&attribute.tokens)? {
                return Ok(Some(attribute));
            }
        }
        Ok(None)
    }
}