From 8f425b9ea725c192a6e42045d9b0d45eec460728 Mon Sep 17 00:00:00 2001 From: Marc Mueller <30130371+cdce8p@users.noreply.github.com> Date: Sat, 6 Apr 2024 21:35:42 +0200 Subject: [PATCH] Improve generic event typing [recorder] (#114736) --- .../components/recorder/entity_registry.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/homeassistant/components/recorder/entity_registry.py b/homeassistant/components/recorder/entity_registry.py index 5bf1856316a..1c0299fc8da 100644 --- a/homeassistant/components/recorder/entity_registry.py +++ b/homeassistant/components/recorder/entity_registry.py @@ -1,8 +1,7 @@ """Recorder entity registry helper.""" -from collections.abc import Mapping import logging -from typing import Any +from typing import TYPE_CHECKING from homeassistant.core import Event, HomeAssistant, callback from homeassistant.helpers import entity_registry as er @@ -19,10 +18,14 @@ def async_setup(hass: HomeAssistant) -> None: """Set up the entity hooks.""" @callback - def _async_entity_id_changed(event: Event) -> None: + def _async_entity_id_changed( + event: Event[er.EventEntityRegistryUpdatedData], + ) -> None: instance = get_instance(hass) - old_entity_id: str = event.data["old_entity_id"] - new_entity_id: str = event.data["entity_id"] + if TYPE_CHECKING: + assert event.data["action"] == "update" and "old_entity_id" in event.data + old_entity_id = event.data["old_entity_id"] + new_entity_id = event.data["entity_id"] instance.async_update_statistics_metadata( old_entity_id, new_statistic_id=new_entity_id ) @@ -31,7 +34,9 @@ def async_setup(hass: HomeAssistant) -> None: ) @callback - def entity_registry_changed_filter(event_data: Mapping[str, Any]) -> bool: + def entity_registry_changed_filter( + event_data: er.EventEntityRegistryUpdatedData, + ) -> bool: """Handle entity_id changed filter.""" return event_data["action"] == "update" and "old_entity_id" in event_data