diff --git a/src/home_assistant/event/context/context.rs b/src/home_assistant/event/context/context.rs
index b033918..2018adc 100644
--- a/src/home_assistant/event/context/context.rs
+++ b/src/home_assistant/event/context/context.rs
@@ -1,24 +1,14 @@
 use pyo3::prelude::*;
 
-use crate::home_assistant::event::event::Event;
-
 use super::id::Id;
 
 /// The context that triggered something.
 #[derive(Debug, FromPyObject)]
-pub struct Context {
+pub struct Context<Event> {
     pub id: Id,
     pub user_id: Option<String>,
     pub parent_id: Option<String>,
-    /// In order to prevent cycles, the user must extract this to an [`Event<Arbitrary>`](super::event::Event) themself (or even specify a specific type parameter!)
-    origin_event: Py<PyAny>,
-}
-
-impl Context {
-    pub fn origin_event<'py, Type: FromPyObject<'py>, Data: FromPyObject<'py>>(
-        &self,
-        py: Python<'py>,
-    ) -> PyResult<Event<Type, Data>> {
-        self.origin_event.extract(py)
-    }
+    /// In order to prevent cycles, the user must decide to pass [`Py<PyAny>`] for the `Event` type here
+    /// or for the `Context` type in [`Event`]
+    pub origin_event: Event,
 }
diff --git a/src/home_assistant/event/event.rs b/src/home_assistant/event/event.rs
index c65a9e1..bf4f62b 100644
--- a/src/home_assistant/event/event.rs
+++ b/src/home_assistant/event/event.rs
@@ -5,20 +5,17 @@ use super::{context::context::Context, event_origin::EventOrigin};
 
 /// Representation of an event within the bus.
 #[derive(Debug, FromPyObject)]
-pub struct Event<Type, Data> {
+pub struct Event<Type, Data, Context> {
     pub event_type: Type,
     pub data: Data,
     pub origin: EventOrigin,
-    /// In order to prevent cycles, the user must extract this to a [`Context`](super::context::Context) themself, using the [`context`](Self::context) method
-    context: Py<PyAny>,
+    /// In order to prevent cycles, the user must decide to pass [`Py<PyAny>`] for the `Context` type here
+    /// or for the `Event` type in [`Context`]
+    pub context: Context,
     time_fired_timestamp: f64,
 }
 
-impl<Type, Data> Event<Type, Data> {
-    pub fn context<'py>(&self, py: Python<'py>) -> PyResult<Context> {
-        self.context.extract(py)
-    }
-
+impl<Type, Data, Context> Event<Type, Data, Context> {
     pub fn time_fired(&self) -> Option<DateTime<Utc>> {
         const NANOS_PER_SEC: i32 = 1_000_000_000;
 
diff --git a/src/home_assistant/event/specific/state_changed.rs b/src/home_assistant/event/specific/state_changed.rs
index 20e152f..9f4a5f9 100644
--- a/src/home_assistant/event/specific/state_changed.rs
+++ b/src/home_assistant/event/specific/state_changed.rs
@@ -24,11 +24,12 @@ impl<'py> FromPyObject<'py> for Type {
 
 #[derive(Debug, FromPyObject)]
 #[pyo3(from_item_all)]
-pub struct Data {
+pub struct Data<OldStateContextEvent, NewStateContextEvent> {
     pub entity_id: EntityId,
-    pub old_state: Option<State>,
-    pub new_state: Option<State>,
+    pub old_state: Option<State<OldStateContextEvent>>,
+    pub new_state: Option<State<NewStateContextEvent>>,
 }
 
 /// A state changed event is fired when on state write the state is changed.
-pub type Event = super::super::event::Event<Type, Data>;
+pub type Event<OldStateContextEvent, NewStateContextEvent, Context> =
+    super::super::event::Event<Type, Data<OldStateContextEvent, NewStateContextEvent>, Context>;
diff --git a/src/home_assistant/state.rs b/src/home_assistant/state.rs
index 59a3e9b..d33d696 100644
--- a/src/home_assistant/state.rs
+++ b/src/home_assistant/state.rs
@@ -6,12 +6,12 @@ use crate::{arbitrary::map::Map, home_assistant::entity_id::EntityId};
 use super::event::context::context::Context;
 
 #[derive(Debug, FromPyObject)]
-pub struct State {
+pub struct State<ContextEvent> {
     pub entity_id: EntityId,
     pub state: String,
     pub attributes: Map,
     pub last_changed: Option<DateTime<Utc>>,
     pub last_reported: Option<DateTime<Utc>>,
     pub last_updated: Option<DateTime<Utc>>,
-    pub context: Context,
+    pub context: Context<ContextEvent>,
 }
diff --git a/src/home_assistant/state_machine.rs b/src/home_assistant/state_machine.rs
index 791af1e..43be851 100644
--- a/src/home_assistant/state_machine.rs
+++ b/src/home_assistant/state_machine.rs
@@ -21,7 +21,11 @@ impl<'py> FromPyObject<'py> for StateMachine {
 }
 
 impl StateMachine {
-    pub fn get(&self, py: &Python, entity_id: EntityId) -> Result<Option<State>, PyErr> {
+    pub fn get<ContextEvent: for<'py> FromPyObject<'py>>(
+        &self,
+        py: &Python,
+        entity_id: EntityId,
+    ) -> Result<Option<State<ContextEvent>>, PyErr> {
         let args = (entity_id.to_string(),);
         let state = self.0.call_method1(*py, "get", args)?;
         state.extract(*py)