From d1d2ba29cee7c566d0f8d1990567697a0ee4465f Mon Sep 17 00:00:00 2001
From: AlephCubed <jothebro500@gmail.com>
Date: Fri, 25 Apr 2025 18:26:13 -0700
Subject: [PATCH] Improved macro error handling.

---
 immediate_stats_macros/src/bevy/butler.rs | 16 +++++++---------
 immediate_stats_macros/src/lib.rs         | 13 ++++++++-----
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/immediate_stats_macros/src/bevy/butler.rs b/immediate_stats_macros/src/bevy/butler.rs
index 7dbc9a8..10f797a 100644
--- a/immediate_stats_macros/src/bevy/butler.rs
+++ b/immediate_stats_macros/src/bevy/butler.rs
@@ -1,28 +1,26 @@
 use darling::ast::NestedMeta;
 use darling::{Error, FromMeta};
-use proc_macro::{self, TokenStream};
-use proc_macro2::Ident;
+use proc_macro2::{Ident, TokenStream};
 use quote::{ToTokens, format_ident, quote};
-use syn::{DeriveInput, Expr, Meta, Path, parse_macro_input};
+use syn::{DeriveInput, Expr, Meta, Path};
 
 // Todo Fix error handling and add documentation.
-pub fn register_systems(input: TokenStream) -> TokenStream {
-    let input = parse_macro_input!(input as DeriveInput);
+pub fn register_systems(input: DeriveInput) -> darling::Result<TokenStream> {
     let struct_name = &input.ident;
 
     let mut butler_attributes = ButlerAttributes::new(struct_name);
 
     for attr in input.attrs {
         if attr.path().is_ident("add_component") {
-            let plugin = PluginPath::from_meta(&attr.meta).unwrap();
+            let plugin = PluginPath::from_meta(&attr.meta)?;
             butler_attributes.component_plugin = Some(plugin);
         } else if attr.path().is_ident("add_resource") {
-            let plugin = PluginPath::from_meta(&attr.meta).unwrap();
+            let plugin = PluginPath::from_meta(&attr.meta)?;
             butler_attributes.resource_plugin = Some(plugin);
         }
     }
 
-    butler_attributes.into_token_stream().into()
+    Ok(butler_attributes.into_token_stream())
 }
 
 pub struct ButlerAttributes<'a> {
@@ -42,7 +40,7 @@ impl<'a> ButlerAttributes<'a> {
 }
 
 impl<'a> ToTokens for ButlerAttributes<'a> {
-    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+    fn to_tokens(&self, tokens: &mut TokenStream) {
         if let Some(plugin_path) = &self.component_plugin {
             let ident = &self.ident;
             let plugin = &plugin_path.0;
diff --git a/immediate_stats_macros/src/lib.rs b/immediate_stats_macros/src/lib.rs
index 8e60dc7..aef4a35 100644
--- a/immediate_stats_macros/src/lib.rs
+++ b/immediate_stats_macros/src/lib.rs
@@ -1,13 +1,12 @@
 #[cfg(feature = "bevy")]
 mod bevy;
 
-use crate::bevy::butler::register_systems;
 use proc_macro_error::{
     emit_call_site_error, emit_call_site_warning, emit_warning, proc_macro_error,
 };
 use proc_macro2::{Span, TokenStream};
 use quote::{ToTokens, quote};
-use syn::{Data, DataEnum, DataStruct, DeriveInput, Field, Ident, Index};
+use syn::{Data, DataEnum, DataStruct, DeriveInput, Field, Ident, Index, parse_macro_input};
 
 #[proc_macro_derive(
     StatContainer,
@@ -15,7 +14,7 @@ use syn::{Data, DataEnum, DataStruct, DeriveInput, Field, Ident, Index};
 )]
 #[proc_macro_error]
 pub fn stat_container_derive(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
-    let tree: DeriveInput = syn::parse(item.clone()).expect("TokenStream must be valid.");
+    let tree: DeriveInput = parse_macro_input!(item as DeriveInput);
 
     let struct_name = &tree.ident;
 
@@ -36,8 +35,12 @@ pub fn stat_container_derive(item: proc_macro::TokenStream) -> proc_macro::Token
 
     #[cfg(feature = "bevy_butler")]
     {
-        let systems: TokenStream = register_systems(item).into();
-        quote! { #result #systems }.into()
+        let systems = bevy::butler::register_systems(tree);
+
+        match systems {
+            Ok(systems) => quote! { #result #systems }.into(),
+            Err(e) => e.write_errors().into(),
+        }
     }
 
     #[cfg(not(feature = "bevy_butler"))]
-- 
2.51.2