From 062434532240ed5596b32eef7fa50ca5e8ddd26b Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Mon, 24 Jul 2023 09:14:10 +0200
Subject: [PATCH] Improve `async_track_entity_registry_updated_event` callback
 typing (#97124)

---
 homeassistant/components/mqtt/mixins.py       |  3 ++-
 .../components/switch_as_x/__init__.py        |  7 +++++--
 homeassistant/helpers/entity.py               |  6 ++++--
 homeassistant/helpers/entity_registry.py      | 21 +++++++++++++++----
 homeassistant/helpers/event.py                |  2 +-
 5 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/homeassistant/components/mqtt/mixins.py b/homeassistant/components/mqtt/mixins.py
index 54dea780dab..0a2ee68f7c4 100644
--- a/homeassistant/components/mqtt/mixins.py
+++ b/homeassistant/components/mqtt/mixins.py
@@ -54,6 +54,7 @@ from homeassistant.helpers.typing import (
     UNDEFINED,
     ConfigType,
     DiscoveryInfoType,
+    EventType,
     UndefinedType,
 )
 from homeassistant.util.json import json_loads
@@ -616,7 +617,7 @@ async def async_remove_discovery_payload(
 async def async_clear_discovery_topic_if_entity_removed(
     hass: HomeAssistant,
     discovery_data: DiscoveryInfoType,
-    event: Event,
+    event: EventType[er.EventEntityRegistryUpdatedData],
 ) -> None:
     """Clear the discovery topic if the entity is removed."""
     if event.data["action"] == "remove":
diff --git a/homeassistant/components/switch_as_x/__init__.py b/homeassistant/components/switch_as_x/__init__.py
index ef64a86c6e8..e2ad91e990e 100644
--- a/homeassistant/components/switch_as_x/__init__.py
+++ b/homeassistant/components/switch_as_x/__init__.py
@@ -8,9 +8,10 @@ import voluptuous as vol
 from homeassistant.components.homeassistant import exposed_entities
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_ENTITY_ID
-from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.event import async_track_entity_registry_updated_event
+from homeassistant.helpers.typing import EventType
 
 from .const import CONF_TARGET_DOMAIN
 from .light import LightSwitch
@@ -55,7 +56,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         )
         return False
 
-    async def async_registry_updated(event: Event) -> None:
+    async def async_registry_updated(
+        event: EventType[er.EventEntityRegistryUpdatedData],
+    ) -> None:
         """Handle entity registry update."""
         data = event.data
         if data["action"] == "remove":
diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py
index 55dc69540fd..a720c1831d7 100644
--- a/homeassistant/helpers/entity.py
+++ b/homeassistant/helpers/entity.py
@@ -45,7 +45,7 @@ from .event import (
     async_track_device_registry_updated_event,
     async_track_entity_registry_updated_event,
 )
-from .typing import UNDEFINED, StateType, UndefinedType
+from .typing import UNDEFINED, EventType, StateType, UndefinedType
 
 if TYPE_CHECKING:
     from .entity_platform import EntityPlatform
@@ -1097,7 +1097,9 @@ class Entity(ABC):
         if self.platform:
             self.hass.data[DATA_ENTITY_SOURCE].pop(self.entity_id)
 
-    async def _async_registry_updated(self, event: Event) -> None:
+    async def _async_registry_updated(
+        self, event: EventType[er.EventEntityRegistryUpdatedData]
+    ) -> None:
         """Handle entity registry update."""
         data = event.data
         if data["action"] == "remove":
diff --git a/homeassistant/helpers/entity_registry.py b/homeassistant/helpers/entity_registry.py
index 248db9d5180..a46dd3c3a52 100644
--- a/homeassistant/helpers/entity_registry.py
+++ b/homeassistant/helpers/entity_registry.py
@@ -108,15 +108,28 @@ class RegistryEntryHider(StrEnum):
     USER = "user"
 
 
-class EventEntityRegistryUpdatedData(TypedDict):
-    """EventEntityRegistryUpdated data."""
+class _EventEntityRegistryUpdatedData_CreateRemove(TypedDict):
+    """EventEntityRegistryUpdated data for action type 'create' and 'remove'."""
 
-    action: Literal["create", "remove", "update"]
+    action: Literal["create", "remove"]
     entity_id: str
-    changes: NotRequired[dict[str, Any]]
+
+
+class _EventEntityRegistryUpdatedData_Update(TypedDict):
+    """EventEntityRegistryUpdated data for action type 'update'."""
+
+    action: Literal["update"]
+    entity_id: str
+    changes: dict[str, Any]  # Required with action == "update"
     old_entity_id: NotRequired[str]
 
 
+EventEntityRegistryUpdatedData = (
+    _EventEntityRegistryUpdatedData_CreateRemove
+    | _EventEntityRegistryUpdatedData_Update
+)
+
+
 EntityOptionsType = Mapping[str, Mapping[str, Any]]
 ReadOnlyEntityOptionsType = ReadOnlyDict[str, Mapping[str, Any]]
 
diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index 12cf58eaa2b..b31efac92bc 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -416,7 +416,7 @@ def _async_dispatch_old_entity_id_or_entity_id_event(
 ) -> None:
     """Dispatch to listeners."""
     if not (
-        callbacks_list := callbacks.get(
+        callbacks_list := callbacks.get(  # type: ignore[call-overload]  # mypy bug?
             event.data.get("old_entity_id", event.data["entity_id"])
         )
     ):