From 1d35c745bb2390ae52dc683991906ba57d4a5bba Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 22 Aug 2024 14:00:30 -0500
Subject: [PATCH] Reduce ESPHome entity state write overhead (#124329)

---
 homeassistant/components/esphome/entity.py     | 8 +++-----
 homeassistant/components/esphome/entry_data.py | 5 ++++-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/homeassistant/components/esphome/entity.py b/homeassistant/components/esphome/entity.py
index 764390c861b..a484a988ea9 100644
--- a/homeassistant/components/esphome/entity.py
+++ b/homeassistant/components/esphome/entity.py
@@ -94,7 +94,6 @@ async def platform_async_setup_entry(
     """
     entry_data = entry.runtime_data
     entry_data.info[info_type] = {}
-    entry_data.state.setdefault(state_type, {})
     platform = entity_platform.async_get_current_platform()
     on_static_info_update = functools.partial(
         async_static_info_updated,
@@ -188,6 +187,7 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
     ) -> None:
         """Initialize."""
         self._entry_data = entry_data
+        self._states = cast(dict[int, _StateT], entry_data.state[state_type])
         assert entry_data.device_info is not None
         device_info = entry_data.device_info
         self._device_info = device_info
@@ -265,11 +265,9 @@ class EsphomeEntity(Entity, Generic[_InfoT, _StateT]):
     @callback
     def _update_state_from_entry_data(self) -> None:
         """Update state from entry data."""
-        state = self._entry_data.state
         key = self._key
-        state_type = self._state_type
-        if has_state := key in state[state_type]:
-            self._state = cast(_StateT, state[state_type][key])
+        if has_state := key in self._states:
+            self._state = self._states[key]
         self._has_state = has_state
 
     @callback
diff --git a/homeassistant/components/esphome/entry_data.py b/homeassistant/components/esphome/entry_data.py
index ff6f048eba1..6fc40612c48 100644
--- a/homeassistant/components/esphome/entry_data.py
+++ b/homeassistant/components/esphome/entry_data.py
@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+from collections import defaultdict
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass, field
 from functools import partial
@@ -111,7 +112,9 @@ class RuntimeEntryData:
     title: str
     client: APIClient
     store: ESPHomeStorage
-    state: dict[type[EntityState], dict[int, EntityState]] = field(default_factory=dict)
+    state: defaultdict[type[EntityState], dict[int, EntityState]] = field(
+        default_factory=lambda: defaultdict(dict)
+    )
     # When the disconnect callback is called, we mark all states
     # as stale so we will always dispatch a state update when the
     # device reconnects. This is the same format as state_subscriptions.