From 006a56130736da6ec33d121f3f0111b089fcb6fe Mon Sep 17 00:00:00 2001
From: Dreaded_X <tim@huizinga.dev>
Date: Sat, 11 Oct 2025 02:05:35 +0200
Subject: [PATCH] feat: Use PartialUserData on proxy type to add trait methods

---
 automation_lib/src/lua/traits.rs | 69 ++++++++++++++++++--------------
 automation_macro/src/device.rs   | 25 +++---------
 2 files changed, 45 insertions(+), 49 deletions(-)

diff --git a/automation_lib/src/lua/traits.rs b/automation_lib/src/lua/traits.rs
index 38c4cb1..4d87728 100644
--- a/automation_lib/src/lua/traits.rs
+++ b/automation_lib/src/lua/traits.rs
@@ -2,21 +2,28 @@ use std::ops::Deref;
 
 // TODO: Enable and disable functions based on query_only and command_only
 
-pub trait Device {
-    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M)
-    where
-        Self: Sized + crate::device::Device + 'static,
-    {
+pub trait PartialUserData<T> {
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M);
+}
+
+pub struct Device;
+
+impl<T> PartialUserData<T> for Device
+where
+    T: crate::device::Device + 'static,
+{
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M) {
         methods.add_async_method("get_id", async |_lua, this, _: ()| Ok(this.get_id()));
     }
 }
-impl<T> Device for T where T: crate::device::Device {}
 
-pub trait OnOff {
-    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M)
-    where
-        Self: Sized + google_home::traits::OnOff + 'static,
-    {
+pub struct OnOff;
+
+impl<T> PartialUserData<T> for OnOff
+where
+    T: google_home::traits::OnOff + 'static,
+{
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M) {
         methods.add_async_method("set_on", async |_lua, this, on: bool| {
             this.deref().set_on(on).await.unwrap();
 
@@ -28,13 +35,14 @@ pub trait OnOff {
         });
     }
 }
-impl<T> OnOff for T where T: google_home::traits::OnOff {}
 
-pub trait Brightness {
-    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M)
-    where
-        Self: Sized + google_home::traits::Brightness + 'static,
-    {
+pub struct Brightness;
+
+impl<T> PartialUserData<T> for Brightness
+where
+    T: google_home::traits::Brightness + 'static,
+{
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M) {
         methods.add_async_method("set_brightness", async |_lua, this, brightness: u8| {
             this.set_brightness(brightness).await.unwrap();
 
@@ -46,13 +54,14 @@ pub trait Brightness {
         });
     }
 }
-impl<T> Brightness for T where T: google_home::traits::Brightness {}
 
-pub trait ColorSetting {
-    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M)
-    where
-        Self: Sized + google_home::traits::ColorSetting + 'static,
-    {
+pub struct ColorSetting;
+
+impl<T> PartialUserData<T> for ColorSetting
+where
+    T: google_home::traits::ColorSetting + 'static,
+{
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M) {
         methods.add_async_method(
             "set_color_temperature",
             async |_lua, this, temperature: u32| {
@@ -69,13 +78,14 @@ pub trait ColorSetting {
         });
     }
 }
-impl<T> ColorSetting for T where T: google_home::traits::ColorSetting {}
 
-pub trait OpenClose {
-    fn add_methods<M: mlua::UserDataMethods<Self>>(methods: &mut M)
-    where
-        Self: Sized + google_home::traits::OpenClose + 'static,
-    {
+pub struct OpenClose;
+
+impl<T> PartialUserData<T> for OpenClose
+where
+    T: google_home::traits::OpenClose + 'static,
+{
+    fn add_methods<M: mlua::UserDataMethods<T>>(methods: &mut M) {
         methods.add_async_method("set_open_percent", async |_lua, this, open_percent: u8| {
             this.set_open_percent(open_percent).await.unwrap();
 
@@ -87,4 +97,3 @@ pub trait OpenClose {
         });
     }
 }
-impl<T> OpenClose for T where T: google_home::traits::OpenClose {}
diff --git a/automation_macro/src/device.rs b/automation_macro/src/device.rs
index f541552..64fcd71 100644
--- a/automation_macro/src/device.rs
+++ b/automation_macro/src/device.rs
@@ -66,18 +66,6 @@ impl Parse for Traits {
     }
 }
 
-impl ToTokens for Traits {
-    fn to_tokens(&self, tokens: &mut TokenStream2) {
-        let Self(traits) = &self;
-
-        tokens.extend(quote! {
-            #(
-                ::automation_lib::lua::traits::#traits::add_methods(methods);
-            )*
-        });
-    }
-}
-
 #[derive(Default)]
 struct Aliases(Vec<syn::Ident>);
 
@@ -138,12 +126,9 @@ impl quote::ToTokens for Implementation {
             traits,
             add_methods,
         } = &self;
+        let Traits(traits) = traits;
 
-        let interfaces: String = traits
-            .0
-            .iter()
-            .map(|tr| format!(", Interface{tr}"))
-            .collect();
+        let interfaces: String = traits.iter().map(|tr| format!(", Interface{tr}")).collect();
 
         tokens.extend(quote! {
             impl mlua::UserData for #name {
@@ -161,9 +146,11 @@ impl quote::ToTokens for Implementation {
                         Ok(b)
                     });
 
-                    ::automation_lib::lua::traits::Device::add_methods(methods);
+                    <::automation_lib::lua::traits::Device as ::automation_lib::lua::traits::PartialUserData<#name>>::add_methods(methods);
 
-                    #traits
+                    #(
+                        <::automation_lib::lua::traits::#traits as ::automation_lib::lua::traits::PartialUserData<#name>>::add_methods(methods);
+                    )*
 
                     #(
                         #add_methods(methods);