Converted macro to derive macro

automation_macro/src/impl_device.rs

@@ -1,70 +1,88 @@
-use proc_macro::TokenStream;
-use quote::quote;
+use proc_macro2::TokenStream;
+use quote::{ToTokens, quote};
 use syn::parse::Parse;
 use syn::punctuated::Punctuated;
-use syn::{Path, Token, parse_macro_input};
+use syn::{AngleBracketedGenericArguments, Attribute, DeriveInput, Ident, Path, Token};
 
-struct ImplDevice {
-    ty: Path,
-    impls: Option<Punctuated<Path, Token![,]>>,
+#[derive(Debug, Default)]
+struct Impl {
+    generics: Option<AngleBracketedGenericArguments>,
+    traits: Vec<Path>,
 }
 
-impl Parse for ImplDevice {
+impl Parse for Impl {
     fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
-        let ty = input.parse()?;
-        let impls = if input.peek(Token![->]) {
-            input.parse::<Token![->]>()?;
-            Some(input.parse_terminated(Path::parse, Token![,])?)
+        let generics = if input.peek(Token![<]) {
+            let generics = input.parse()?;
+            input.parse::<Token![:]>()?;
+
+            Some(generics)
         } else {
             None
         };
 
-        Ok(ImplDevice { ty, impls })
+        let traits: Punctuated<_, _> = input.parse_terminated(Path::parse, Token![,])?;
+        let traits = traits.into_iter().collect();
+
+        Ok(Impl { generics, traits })
     }
 }
 
-pub fn impl_device_macro(input: proc_macro::TokenStream) -> TokenStream {
-    let ImplDevice { ty, impls } = parse_macro_input!(input as ImplDevice);
+impl Impl {
+    fn generate(&self, name: &Ident) -> TokenStream {
+        let generics = &self.generics;
 
-    let impls: Vec<_> = impls
-        .iter()
-        .flatten()
-        .map(|i| {
-            let ident = i
-                .segments
-                .last()
-                .expect("There should be at least one segment")
-                .ident
-                .clone();
-
-            quote! {
-                ::automation_lib::lua::traits::#ident::add_methods(methods);
+        // If an identifier is specified, assume it is placed in ::automation_lib::lua::traits,
+        // otherwise use the provided path
+        let traits = self.traits.iter().map(|t| {
+            if let Some(ident) = t.get_ident() {
+                quote! {::automation_lib::lua::traits::#ident }
+            } else {
+                t.to_token_stream()
             }
-        })
-        .collect();
+        });
 
-    quote! {
-        impl mlua::UserData for #ty {
-            fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
-                methods.add_async_function("new", |_lua, config| async {
-                    let device: #ty = LuaDeviceCreate::create(config)
-                        .await
-                        .map_err(mlua::ExternalError::into_lua_err)?;
+        quote! {
+            impl mlua::UserData for #name #generics {
+                fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M) {
+                    methods.add_async_function("new", |_lua, config| async {
+                        let device: Self = LuaDeviceCreate::create(config)
+                            .await
+                            .map_err(mlua::ExternalError::into_lua_err)?;
 
-                    Ok(device)
-                });
+                        Ok(device)
+                    });
 
-                methods.add_method("__box", |_lua, this, _: ()| {
-                    let b: Box<dyn Device> = Box::new(this.clone());
-                    Ok(b)
-                });
+                    methods.add_method("__box", |_lua, this, _: ()| {
+                        let b: Box<dyn Device> = Box::new(this.clone());
+                        Ok(b)
+                    });
 
-                methods.add_async_method("get_id", |_lua, this, _: ()| async move { Ok(this.get_id()) });
+                    methods.add_async_method("get_id", |_lua, this, _: ()| async move { Ok(this.get_id()) });
 
-                #(
-                    #impls
-                )*
+                    #(
+                        #traits::add_methods(methods);
+                    )*
+                }
             }
         }
-    }.into()
+    }
+}
+
+pub fn impl_device_macro(ast: &DeriveInput) -> TokenStream {
+    let name = &ast.ident;
+
+    let impls: TokenStream = ast
+        .attrs
+        .iter()
+        .filter(|attr| attr.path().is_ident("traits"))
+        .flat_map(Attribute::parse_args::<Impl>)
+        .map(|im| im.generate(name))
+        .collect();
+
+    if impls.is_empty() {
+        Impl::default().generate(name)
+    } else {
+        impls
+    }
 }

automation_macro/src/lib.rs

@@ -14,7 +14,9 @@ pub fn lua_device_config_derive(input: proc_macro::TokenStream) -> proc_macro::T
     impl_lua_device_config_macro(&ast).into()
 }
 
-#[proc_macro]
+#[proc_macro_derive(LuaDevice, attributes(traits))]
 pub fn impl_device(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
-    impl_device_macro(input)
+    let ast = parse_macro_input!(input as DeriveInput);
+
+    impl_device_macro(&ast).into()
 }