From 28ad5d7345d85ef652e8eec4a8b68f7bf6647d79 Mon Sep 17 00:00:00 2001
From: Jacob <jacob@katniss.top>
Date: Thu, 13 Mar 2025 16:15:46 -0400
Subject: [PATCH] feat: static `Attributes` type

---
 src/home_assistant/event/specific/state_changed.rs | 14 +++++++++-----
 src/home_assistant/state.rs                        |  4 ++--
 src/home_assistant/state_machine.rs                |  4 ++--
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/src/home_assistant/event/specific/state_changed.rs b/src/home_assistant/event/specific/state_changed.rs
index 9f4a5f9..45ad34c 100644
--- a/src/home_assistant/event/specific/state_changed.rs
+++ b/src/home_assistant/event/specific/state_changed.rs
@@ -24,12 +24,16 @@ impl<'py> FromPyObject<'py> for Type {
 
 #[derive(Debug, FromPyObject)]
 #[pyo3(from_item_all)]
-pub struct Data<OldStateContextEvent, NewStateContextEvent> {
+pub struct Data<OldAttributes, OldStateContextEvent, NewAttributes, NewStateContextEvent> {
     pub entity_id: EntityId,
-    pub old_state: Option<State<OldStateContextEvent>>,
-    pub new_state: Option<State<NewStateContextEvent>>,
+    pub old_state: Option<State<OldAttributes, OldStateContextEvent>>,
+    pub new_state: Option<State<NewAttributes, NewStateContextEvent>>,
 }
 
 /// A state changed event is fired when on state write the state is changed.
-pub type Event<OldStateContextEvent, NewStateContextEvent, Context> =
-    super::super::event::Event<Type, Data<OldStateContextEvent, NewStateContextEvent>, Context>;
+pub type Event<OldAttributes, OldStateContextEvent, NewAttributes, NewStateContextEvent, Context> =
+    super::super::event::Event<
+        Type,
+        Data<OldAttributes, OldStateContextEvent, NewAttributes, NewStateContextEvent>,
+        Context,
+    >;
diff --git a/src/home_assistant/state.rs b/src/home_assistant/state.rs
index d33d696..4f48996 100644
--- a/src/home_assistant/state.rs
+++ b/src/home_assistant/state.rs
@@ -6,10 +6,10 @@ use crate::{arbitrary::map::Map, home_assistant::entity_id::EntityId};
 use super::event::context::context::Context;
 
 #[derive(Debug, FromPyObject)]
-pub struct State<ContextEvent> {
+pub struct State<Attributes, ContextEvent> {
     pub entity_id: EntityId,
     pub state: String,
-    pub attributes: Map,
+    pub attributes: Attributes,
     pub last_changed: Option<DateTime<Utc>>,
     pub last_reported: Option<DateTime<Utc>>,
     pub last_updated: Option<DateTime<Utc>>,
diff --git a/src/home_assistant/state_machine.rs b/src/home_assistant/state_machine.rs
index 43be851..77adf58 100644
--- a/src/home_assistant/state_machine.rs
+++ b/src/home_assistant/state_machine.rs
@@ -21,11 +21,11 @@ impl<'py> FromPyObject<'py> for StateMachine {
 }
 
 impl StateMachine {
-    pub fn get<ContextEvent: for<'py> FromPyObject<'py>>(
+    pub fn get<Attributes: for<'py> FromPyObject<'py>, ContextEvent: for<'py> FromPyObject<'py>>(
         &self,
         py: &Python,
         entity_id: EntityId,
-    ) -> Result<Option<State<ContextEvent>>, PyErr> {
+    ) -> Result<Option<State<Attributes, ContextEvent>>, PyErr> {
         let args = (entity_id.to_string(),);
         let state = self.0.call_method1(*py, "get", args)?;
         state.extract(*py)