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}