From 0cedf8949dc682a1f91510d2c373a2a29aca62f6 Mon Sep 17 00:00:00 2001
From: Victor Koenders <gpg@trangar.com>
Date: Sat, 4 Dec 2021 12:04:58 +0100
Subject: [PATCH] Added a way to relatively simply add custom attributes

---
 src/lib.rs              |  4 ++--
 src/parse/attributes.rs |  9 +++++++++
 src/parse/body.rs       | 17 ++++++++++++++++-
 src/parse/mod.rs        |  2 +-
 4 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 3b273d0..7e8354d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -80,7 +80,7 @@ pub use self::error::Error;
 /// Useful includes
 pub mod prelude {
     pub use crate::generate::FnSelfArg;
-    pub use crate::parse::{Body, Parse};
+    pub use crate::parse::{Body, FromAttribute, Parse};
     pub use crate::Result;
     pub use proc_macro2::*;
 }
@@ -90,7 +90,7 @@ pub mod prelude {
     extern crate proc_macro;
 
     pub use crate::generate::FnSelfArg;
-    pub use crate::parse::{Body, Parse};
+    pub use crate::parse::{Body, FromAttribute, Parse};
     pub use crate::Result;
     pub use proc_macro::*;
 }
diff --git a/src/parse/attributes.rs b/src/parse/attributes.rs
index 9283a2f..cb88fd4 100644
--- a/src/parse/attributes.rs
+++ b/src/parse/attributes.rs
@@ -95,3 +95,12 @@ fn test_attributes_try_take() {
         x => panic!("Expected ident, found {:?}", x),
     }
 }
+
+/// Helper trait for functions like:
+/// - [`IdentOrIndex::has_field_attribute`]
+///
+/// This can be implemented on your own type to make parsing easier.
+pub trait FromAttribute: Sized {
+    /// Try to parse the given group into your own type. Return `None` if the parsing failed or if the attribute was not this type.
+    fn parse(group: &Group) -> Option<Self>;
+}
diff --git a/src/parse/body.rs b/src/parse/body.rs
index ad87554..4b21f3d 100644
--- a/src/parse/body.rs
+++ b/src/parse/body.rs
@@ -1,5 +1,5 @@
 use super::attributes::AttributeLocation;
-use super::{utils::*, Attribute, Visibility};
+use super::{utils::*, Attribute, FromAttribute, Visibility};
 use crate::prelude::{Delimiter, Ident, Literal, Span, TokenTree};
 use crate::{Error, Result};
 use std::iter::Peekable;
@@ -507,6 +507,21 @@ impl<'a> IdentOrIndex<'a> {
             }
         }
     }
+
+    pub fn has_field_attribute<T: FromAttribute + PartialEq<T>>(&self, attrib: T) -> bool {
+        let attributes = match self {
+            Self::Ident { attributes, .. } => attributes,
+            Self::Index { attributes, .. } => attributes,
+        };
+        for attribute in attributes.iter() {
+            if let Some(attribute) = T::parse(&attribute.tokens) {
+                if attribute == attrib {
+                    return true;
+                }
+            }
+        }
+        false
+    }
 }
 
 impl std::fmt::Display for IdentOrIndex<'_> {
diff --git a/src/parse/mod.rs b/src/parse/mod.rs
index 55e86d0..e4bf3ef 100644
--- a/src/parse/mod.rs
+++ b/src/parse/mod.rs
@@ -9,7 +9,7 @@ mod generics;
 mod utils;
 mod visibility;
 
-pub use self::attributes::{Attribute, AttributeLocation};
+pub use self::attributes::{Attribute, AttributeLocation, FromAttribute};
 pub use self::body::{EnumBody, EnumVariant, Fields, StructBody, UnnamedField};
 pub(crate) use self::data_type::DataType;
 pub use self::generics::{Generic, GenericConstraints, Generics, Lifetime, SimpleGeneric};