Improve async_track_entity_registry_updated_event callback typing (#97124)

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":

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":

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":

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]]
 

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"])
         )
     ):