From 70d1bbb20d42fcb316001004b5a7f92593e7d192 Mon Sep 17 00:00:00 2001
From: Franck Nijhof <git@frenck.dev>
Date: Sun, 18 Feb 2024 14:57:15 +0100
Subject: [PATCH] Improve floor registry event typing (#110844)

---
 homeassistant/helpers/area_registry.py  | 14 ++++++-----
 homeassistant/helpers/floor_registry.py | 31 +++++++++++++++++++++----
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/homeassistant/helpers/area_registry.py b/homeassistant/helpers/area_registry.py
index 256a46cb62872..c0d2efcb99e18 100644
--- a/homeassistant/helpers/area_registry.py
+++ b/homeassistant/helpers/area_registry.py
@@ -6,7 +6,7 @@
 import dataclasses
 from typing import Any, Literal, TypedDict, cast
 
-from homeassistant.core import Event, HomeAssistant, callback
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.util import slugify
 
 from . import device_registry as dr, entity_registry as er
@@ -344,12 +344,14 @@ def _async_setup_cleanup(self) -> None:
         from . import floor_registry as fr  # Circular dependency
 
         @callback
-        def _floor_removed_from_registry_filter(event: Event) -> bool:
+        def _floor_removed_from_registry_filter(
+            event: fr.EventFloorRegistryUpdated,
+        ) -> bool:
             """Filter all except for the remove action from floor registry events."""
-            return bool(event.data["action"] == "remove")
+            return event.data["action"] == "remove"
 
         @callback
-        def _handle_floor_registry_update(event: Event) -> None:
+        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["floor_id"]
             for area_id, area in self.areas.items():
@@ -358,8 +360,8 @@ def _handle_floor_registry_update(event: Event) -> None:
 
         self.hass.bus.async_listen(
             event_type=fr.EVENT_FLOOR_REGISTRY_UPDATED,
-            event_filter=_floor_removed_from_registry_filter,
-            listener=_handle_floor_registry_update,
+            event_filter=_floor_removed_from_registry_filter,  # type: ignore[arg-type]
+            listener=_handle_floor_registry_update,  # type: ignore[arg-type]
         )
 
 
diff --git a/homeassistant/helpers/floor_registry.py b/homeassistant/helpers/floor_registry.py
index 7d2e060f6550b..e90956e81a207 100644
--- a/homeassistant/helpers/floor_registry.py
+++ b/homeassistant/helpers/floor_registry.py
@@ -5,12 +5,12 @@
 from collections.abc import Iterable, ValuesView
 import dataclasses
 from dataclasses import dataclass
-from typing import cast
+from typing import Literal, TypedDict, cast
 
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.util import slugify
 
-from .typing import UNDEFINED, UndefinedType
+from .typing import UNDEFINED, EventType, UndefinedType
 
 DATA_REGISTRY = "floor_registry"
 EVENT_FLOOR_REGISTRY_UPDATED = "floor_registry_updated"
@@ -19,6 +19,16 @@
 SAVE_DELAY = 10
 
 
+class EventFloorRegistryUpdatedData(TypedDict):
+    """Event data for when the floor registry is updated."""
+
+    action: Literal["create", "remove", "update"]
+    floor_id: str
+
+
+EventFloorRegistryUpdated = EventType[EventFloorRegistryUpdatedData]
+
+
 @dataclass(slots=True, kw_only=True, frozen=True)
 class FloorEntry:
     """Floor registry entry."""
@@ -151,7 +161,10 @@ def async_create(
         self.async_schedule_save()
         self.hass.bus.async_fire(
             EVENT_FLOOR_REGISTRY_UPDATED,
-            {"action": "create", "floor_id": floor_id},
+            EventFloorRegistryUpdatedData(
+                action="create",
+                floor_id=floor_id,
+            ),
         )
         return floor
 
@@ -160,7 +173,11 @@ def async_delete(self, floor_id: str) -> None:
         """Delete floor."""
         del self.floors[floor_id]
         self.hass.bus.async_fire(
-            EVENT_FLOOR_REGISTRY_UPDATED, {"action": "remove", "floor_id": floor_id}
+            EVENT_FLOOR_REGISTRY_UPDATED,
+            EventFloorRegistryUpdatedData(
+                action="remove",
+                floor_id=floor_id,
+            ),
         )
         self.async_schedule_save()
 
@@ -196,7 +213,11 @@ def async_update(
 
         self.async_schedule_save()
         self.hass.bus.async_fire(
-            EVENT_FLOOR_REGISTRY_UPDATED, {"action": "update", "floor_id": floor_id}
+            EVENT_FLOOR_REGISTRY_UPDATED,
+            EventFloorRegistryUpdatedData(
+                action="update",
+                floor_id=floor_id,
+            ),
         )
 
         return new