From a165064e9dc60750bb80fc2dd335b7667d59fec9 Mon Sep 17 00:00:00 2001
From: Erik Montnemery <erik@montnemery.com>
Date: Thu, 27 Jun 2024 13:08:19 +0200
Subject: [PATCH] Improve typing of state event helpers (#120639)

---
 homeassistant/core.py          | 15 +++++++++------
 homeassistant/helpers/event.py | 10 ++++++----
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index 2b43b2d40ff..71ee5f4bd1d 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -158,26 +158,29 @@ class ConfigSource(enum.StrEnum):
     YAML = "yaml"
 
 
-class EventStateChangedData(TypedDict):
+class EventStateEventData(TypedDict):
+    """Base class for EVENT_STATE_CHANGED and EVENT_STATE_CHANGED data."""
+
+    entity_id: str
+    new_state: State | None
+
+
+class EventStateChangedData(EventStateEventData):
     """EVENT_STATE_CHANGED data.
 
     A state changed event is fired when on state write when the state is changed.
     """
 
-    entity_id: str
     old_state: State | None
-    new_state: State | None
 
 
-class EventStateReportedData(TypedDict):
+class EventStateReportedData(EventStateEventData):
     """EVENT_STATE_REPORTED data.
 
     A state reported event is fired when on state write when the state is unchanged.
     """
 
-    entity_id: str
     old_last_reported: datetime.datetime
-    new_state: State | None
 
 
 # SOURCE_* are deprecated as of Home Assistant 2022.2, use ConfigSource instead
diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index ebd51948e3b..fa409269ad6 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -27,6 +27,7 @@ from homeassistant.core import (
     Event,
     # Explicit reexport of 'EventStateChangedData' for backwards compatibility
     EventStateChangedData as EventStateChangedData,  # noqa: PLC0414
+    EventStateEventData,
     EventStateReportedData,
     HassJob,
     HassJobType,
@@ -89,6 +90,7 @@ RANDOM_MICROSECOND_MIN = 50000
 RANDOM_MICROSECOND_MAX = 500000
 
 _TypedDictT = TypeVar("_TypedDictT", bound=Mapping[str, Any])
+_StateEventDataT = TypeVar("_StateEventDataT", bound=EventStateEventData)
 
 
 @dataclass(slots=True, frozen=True)
@@ -329,8 +331,8 @@ def async_track_state_change_event(
 @callback
 def _async_dispatch_entity_id_event(
     hass: HomeAssistant,
-    callbacks: dict[str, list[HassJob[[Event[_TypedDictT]], Any]]],
-    event: Event[_TypedDictT],
+    callbacks: dict[str, list[HassJob[[Event[_StateEventDataT]], Any]]],
+    event: Event[_StateEventDataT],
 ) -> None:
     """Dispatch to listeners."""
     if not (callbacks_list := callbacks.get(event.data["entity_id"])):
@@ -349,8 +351,8 @@ def _async_dispatch_entity_id_event(
 @callback
 def _async_state_filter(
     hass: HomeAssistant,
-    callbacks: dict[str, list[HassJob[[Event[_TypedDictT]], Any]]],
-    event_data: _TypedDictT,
+    callbacks: dict[str, list[HassJob[[Event[_StateEventDataT]], Any]]],
+    event_data: _StateEventDataT,
 ) -> bool:
     """Filter state changes by entity_id."""
     return event_data["entity_id"] in callbacks