Simplify esphome state updates (#73409)

homeassistant/components/esphome/__init__.py

@@ -568,6 +568,7 @@ async def platform_async_setup_entry(
     @callback
     def async_list_entities(infos: list[EntityInfo]) -> None:
         """Update entities of this platform when entities are listed."""
+        key_to_component = entry_data.key_to_component
         old_infos = entry_data.info[component_key]
         new_infos: dict[int, EntityInfo] = {}
         add_entities = []
@@ -586,10 +587,12 @@ async def platform_async_setup_entry(
                 entity = entity_type(entry_data, component_key, info.key)
                 add_entities.append(entity)
             new_infos[info.key] = info
+            key_to_component[info.key] = component_key
 
         # Remove old entities
         for info in old_infos.values():
             entry_data.async_remove_entity(hass, component_key, info.key)
+            key_to_component.pop(info.key, None)
 
         # First copy the now-old info into the backup object
         entry_data.old_info[component_key] = entry_data.info[component_key]
@@ -604,22 +607,6 @@ async def platform_async_setup_entry(
         async_dispatcher_connect(hass, signal, async_list_entities)
     )
 
-    @callback
-    def async_entity_state(state: EntityState) -> None:
-        """Notify the appropriate entity of an updated state."""
-        if not isinstance(state, state_type):
-            return
-        # cast back to upper type, otherwise mypy gets confused
-        state = cast(EntityState, state)
-
-        entry_data.state[component_key][state.key] = state
-        entry_data.async_update_entity(hass, component_key, state.key)
-
-    signal = f"esphome_{entry.entry_id}_on_state"
-    entry_data.cleanup_callbacks.append(
-        async_dispatcher_connect(hass, signal, async_entity_state)
-    )
-
 
 _PropT = TypeVar("_PropT", bound=Callable[..., Any])
 

homeassistant/components/esphome/entry_data.py

@@ -65,6 +65,7 @@ class RuntimeEntryData:
     store: Store
     state: dict[str, dict[int, EntityState]] = field(default_factory=dict)
     info: dict[str, dict[int, EntityInfo]] = field(default_factory=dict)
+    key_to_component: dict[int, str] = field(default_factory=dict)
 
     # A second list of EntityInfo objects
     # This is necessary for when an entity is being removed. HA requires
@@ -82,14 +83,6 @@ class RuntimeEntryData:
     platform_load_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
     _storage_contents: dict[str, Any] | None = None
 
-    @callback
-    def async_update_entity(
-        self, hass: HomeAssistant, component_key: str, key: int
-    ) -> None:
-        """Schedule the update of an entity."""
-        signal = f"esphome_{self.entry_id}_update_{component_key}_{key}"
-        async_dispatcher_send(hass, signal)
-
     @callback
     def async_remove_entity(
         self, hass: HomeAssistant, component_key: str, key: int
@@ -131,9 +124,11 @@ class RuntimeEntryData:
 
     @callback
     def async_update_state(self, hass: HomeAssistant, state: EntityState) -> None:
-        """Distribute an update of state information to all platforms."""
-        signal = f"esphome_{self.entry_id}_on_state"
-        async_dispatcher_send(hass, signal, state)
+        """Distribute an update of state information to the target."""
+        component_key = self.key_to_component[state.key]
+        self.state[component_key][state.key] = state
+        signal = f"esphome_{self.entry_id}_update_{component_key}_{state.key}"
+        async_dispatcher_send(hass, signal)
 
     @callback
     def async_update_device_state(self, hass: HomeAssistant) -> None: