Improve floor registry event typing (#110844)

This commit is contained in:
Franck Nijhof 2024-02-18 14:57:15 +01:00 committed by GitHub
parent 1c55ba0cb2
commit 70d1bbb20d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 34 additions and 11 deletions

View file

@ -6,7 +6,7 @@ from collections.abc import Iterable, ValuesView
import dataclasses import dataclasses
from typing import Any, Literal, TypedDict, cast 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 homeassistant.util import slugify
from . import device_registry as dr, entity_registry as er from . import device_registry as dr, entity_registry as er
@ -344,12 +344,14 @@ class AreaRegistry:
from . import floor_registry as fr # Circular dependency from . import floor_registry as fr # Circular dependency
@callback @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.""" """Filter all except for the remove action from floor registry events."""
return bool(event.data["action"] == "remove") return event.data["action"] == "remove"
@callback @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.""" """Update areas that are associated with a floor that has been removed."""
floor_id = event.data["floor_id"] floor_id = event.data["floor_id"]
for area_id, area in self.areas.items(): for area_id, area in self.areas.items():
@ -358,8 +360,8 @@ class AreaRegistry:
self.hass.bus.async_listen( self.hass.bus.async_listen(
event_type=fr.EVENT_FLOOR_REGISTRY_UPDATED, event_type=fr.EVENT_FLOOR_REGISTRY_UPDATED,
event_filter=_floor_removed_from_registry_filter, event_filter=_floor_removed_from_registry_filter, # type: ignore[arg-type]
listener=_handle_floor_registry_update, listener=_handle_floor_registry_update, # type: ignore[arg-type]
) )

View file

@ -5,12 +5,12 @@ from collections import UserDict
from collections.abc import Iterable, ValuesView from collections.abc import Iterable, ValuesView
import dataclasses import dataclasses
from dataclasses import dataclass from dataclasses import dataclass
from typing import cast from typing import Literal, TypedDict, cast
from homeassistant.core import HomeAssistant, callback from homeassistant.core import HomeAssistant, callback
from homeassistant.util import slugify from homeassistant.util import slugify
from .typing import UNDEFINED, UndefinedType from .typing import UNDEFINED, EventType, UndefinedType
DATA_REGISTRY = "floor_registry" DATA_REGISTRY = "floor_registry"
EVENT_FLOOR_REGISTRY_UPDATED = "floor_registry_updated" EVENT_FLOOR_REGISTRY_UPDATED = "floor_registry_updated"
@ -19,6 +19,16 @@ STORAGE_VERSION_MAJOR = 1
SAVE_DELAY = 10 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) @dataclass(slots=True, kw_only=True, frozen=True)
class FloorEntry: class FloorEntry:
"""Floor registry entry.""" """Floor registry entry."""
@ -151,7 +161,10 @@ class FloorRegistry:
self.async_schedule_save() self.async_schedule_save()
self.hass.bus.async_fire( self.hass.bus.async_fire(
EVENT_FLOOR_REGISTRY_UPDATED, EVENT_FLOOR_REGISTRY_UPDATED,
{"action": "create", "floor_id": floor_id}, EventFloorRegistryUpdatedData(
action="create",
floor_id=floor_id,
),
) )
return floor return floor
@ -160,7 +173,11 @@ class FloorRegistry:
"""Delete floor.""" """Delete floor."""
del self.floors[floor_id] del self.floors[floor_id]
self.hass.bus.async_fire( 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() self.async_schedule_save()
@ -196,7 +213,11 @@ class FloorRegistry:
self.async_schedule_save() self.async_schedule_save()
self.hass.bus.async_fire( 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 return new