Improve typing of floor registry events (home-assistant#157624)

Author: emontnemery

homeassistant/helpers/area_registry.py

@@ -516,9 +516,9 @@ def _removed_from_registry_filter(
         @callback
         def _handle_floor_registry_update(event: fr.EventFloorRegistryUpdated) -> None:
             """Update areas that are associated with a floor that has been removed."""
-            floor_id = event.data.get("floor_id")
-            if floor_id is None:
-                return
+            if TYPE_CHECKING:
+                assert event.data["action"] == "remove"
+            floor_id = event.data["floor_id"]
             for area in self.areas.get_areas_for_floor(floor_id):
                 self.async_update(area.id, floor_id=None)
 

homeassistant/helpers/floor_registry.py

@@ -51,12 +51,23 @@ class FloorRegistryStoreData(TypedDict):
     floors: list[_FloorStoreData]
 
 
-class EventFloorRegistryUpdatedData(TypedDict):
+class _EventFloorRegistryUpdatedData_Create_Remove_Update(TypedDict):
     """Event data for when the floor registry is updated."""
 
-    action: Literal["create", "remove", "update", "reorder"]
-    floor_id: str | None
+    action: Literal["create", "remove", "update"]
+    floor_id: str
+
+
+class _EventFloorRegistryUpdatedData_Reorder(TypedDict):
+    """Event data for when the floor registry is updated."""
 
+    action: Literal["reorder"]
+
+
+type EventFloorRegistryUpdatedData = (
+    _EventFloorRegistryUpdatedData_Create_Remove_Update
+    | _EventFloorRegistryUpdatedData_Reorder
+)
 
 type EventFloorRegistryUpdated = Event[EventFloorRegistryUpdatedData]
 
@@ -200,7 +211,9 @@ def async_create(
 
         self.hass.bus.async_fire_internal(
             EVENT_FLOOR_REGISTRY_UPDATED,
-            EventFloorRegistryUpdatedData(action="create", floor_id=floor_id),
+            _EventFloorRegistryUpdatedData_Create_Remove_Update(
+                action="create", floor_id=floor_id
+            ),
         )
         return floor
 
@@ -211,7 +224,7 @@ def async_delete(self, floor_id: str) -> None:
         del self.floors[floor_id]
         self.hass.bus.async_fire_internal(
             EVENT_FLOOR_REGISTRY_UPDATED,
-            EventFloorRegistryUpdatedData(
+            _EventFloorRegistryUpdatedData_Create_Remove_Update(
                 action="remove",
                 floor_id=floor_id,
             ),
@@ -253,7 +266,7 @@ def async_update(
         self.async_schedule_save()
         self.hass.bus.async_fire_internal(
             EVENT_FLOOR_REGISTRY_UPDATED,
-            EventFloorRegistryUpdatedData(
+            _EventFloorRegistryUpdatedData_Create_Remove_Update(
                 action="update",
                 floor_id=floor_id,
             ),
@@ -280,7 +293,7 @@ def async_reorder(self, floor_ids: list[str]) -> None:
         self.async_schedule_save()
         self.hass.bus.async_fire_internal(
             EVENT_FLOOR_REGISTRY_UPDATED,
-            EventFloorRegistryUpdatedData(action="reorder", floor_id=None),
+            _EventFloorRegistryUpdatedData_Reorder(action="reorder"),
         )
 
     async def async_load(self) -> None: