Skip to content

Commit

Permalink
Improve floor registry event typing (#110844)
Browse files Browse the repository at this point in the history
  • Loading branch information
frenck authored Feb 18, 2024
1 parent 1c55ba0 commit 70d1bbb
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 11 deletions.
14 changes: 8 additions & 6 deletions homeassistant/helpers/area_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand All @@ -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]
)


Expand Down
31 changes: 26 additions & 5 deletions homeassistant/helpers/floor_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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."""
Expand Down Expand Up @@ -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

Expand All @@ -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()

Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 70d1bbb

Please sign in to comment.