Migrate ESPHome to use entry.runtime_data (#120402)

homeassistant/components/esphome/__init__.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from aioesphomeapi import APIClient
 
 from homeassistant.components import zeroconf
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PASSWORD,
@@ -21,7 +20,7 @@ from .dashboard import async_setup as async_setup_dashboard
 from .domain_data import DomainData
 
 # Import config flow so that it's added to the registry
-from .entry_data import RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
 from .manager import ESPHomeManager, cleanup_instance
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
@@ -33,7 +32,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     return True
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: ESPHomeConfigEntry) -> bool:
     """Set up the esphome component."""
     host = entry.data[CONF_HOST]
     port = entry.data[CONF_PORT]
@@ -59,7 +58,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         store=domain_data.get_or_create_store(hass, entry),
         original_options=dict(entry.options),
     )
-    domain_data.set_entry_data(entry, entry_data)
+    entry.runtime_data = entry_data
 
     manager = ESPHomeManager(
         hass, entry, host, password, cli, zeroconf_instance, domain_data, entry_data
@@ -69,7 +68,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: ESPHomeConfigEntry) -> bool:
     """Unload an esphome config entry."""
     entry_data = await cleanup_instance(hass, entry)
     return await hass.config_entries.async_unload_platforms(
@@ -77,6 +76,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     )
 
 
-async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
+async def async_remove_entry(hass: HomeAssistant, entry: ESPHomeConfigEntry) -> None:
     """Remove an esphome config entry."""
     await DomainData.get(hass).get_or_create_store(hass, entry).async_remove()

homeassistant/components/esphome/alarm_control_panel.py

@@ -16,7 +16,6 @@ from homeassistant.components.alarm_control_panel import (
     AlarmControlPanelEntityFeature,
     CodeFormat,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     STATE_ALARM_ARMED_AWAY,
     STATE_ALARM_ARMED_CUSTOM_BYPASS,
@@ -38,6 +37,7 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 from .enum_mapper import EsphomeEnumMapper
 
 _ESPHOME_ACP_STATE_TO_HASS_STATE: EsphomeEnumMapper[AlarmControlPanelState, str] = (
@@ -70,7 +70,9 @@ class EspHomeACPFeatures(APIIntEnum):
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome switches based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/binary_sensor.py

@@ -9,17 +9,18 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
 
-from .domain_data import DomainData
 from .entity import EsphomeAssistEntity, EsphomeEntity, platform_async_setup_entry
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome binary sensors based on a config entry."""
     await platform_async_setup_entry(
@@ -31,7 +32,7 @@ async def async_setup_entry(
         state_type=BinarySensorState,
     )
 
-    entry_data = DomainData.get(hass).get_entry_data(entry)
+    entry_data = entry.runtime_data
     assert entry_data.device_info is not None
     if entry_data.device_info.voice_assistant_feature_flags_compat(
         entry_data.api_version

homeassistant/components/esphome/button.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 from aioesphomeapi import ButtonInfo, EntityInfo, EntityState
 
 from homeassistant.components.button import ButtonDeviceClass, ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
@@ -15,10 +14,13 @@ from .entity import (
     convert_api_error_ha_error,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome buttons based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/camera.py

@@ -12,15 +12,17 @@ from aiohttp import web
 
 from homeassistant.components import camera
 from homeassistant.components.camera import Camera
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .entity import EsphomeEntity, platform_async_setup_entry
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up esphome cameras based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/climate.py

@@ -45,7 +45,6 @@ from homeassistant.components.climate import (
     HVACAction,
     HVACMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_TEMPERATURE,
     PRECISION_HALVES,
@@ -62,13 +61,16 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 from .enum_mapper import EsphomeEnumMapper
 
 FAN_QUIET = "quiet"
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome climate devices based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/cover.py

@@ -13,7 +13,6 @@ from homeassistant.components.cover import (
     CoverEntity,
     CoverEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
@@ -24,10 +23,13 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome covers based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/dashboard.py

@@ -106,7 +106,7 @@ class ESPHomeDashboardManager:
         reloads = [
             hass.config_entries.async_reload(entry.entry_id)
             for entry in hass.config_entries.async_entries(DOMAIN)
-            if entry.state == ConfigEntryState.LOADED
+            if entry.state is ConfigEntryState.LOADED
         ]
         # Re-auth flows will check the dashboard for encryption key when the form is requested
         # but we only trigger reauth if the dashboard is available.

homeassistant/components/esphome/date.py

@@ -7,16 +7,16 @@ from datetime import date
 from aioesphomeapi import DateInfo, DateState
 
 from homeassistant.components.date import DateEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .entity import EsphomeEntity, esphome_state_property, platform_async_setup_entry
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: ESPHomeConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up esphome dates based on a config entry."""

homeassistant/components/esphome/datetime.py

@@ -7,17 +7,17 @@ from datetime import datetime
 from aioesphomeapi import DateTimeInfo, DateTimeState
 
 from homeassistant.components.datetime import DateTimeEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 import homeassistant.util.dt as dt_util
 
 from .entity import EsphomeEntity, esphome_state_property, platform_async_setup_entry
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: ESPHomeConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up esphome datetimes based on a config entry."""

homeassistant/components/esphome/diagnostics.py

@@ -6,12 +6,12 @@ from typing import Any
 
 from homeassistant.components.bluetooth import async_scanner_by_source
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_PASSWORD
 from homeassistant.core import HomeAssistant
 
-from . import CONF_NOISE_PSK, DomainData
+from . import CONF_NOISE_PSK
 from .dashboard import async_get_dashboard
+from .entry_data import ESPHomeConfigEntry
 
 CONF_MAC_ADDRESS = "mac_address"
 
@@ -19,14 +19,14 @@ REDACT_KEYS = {CONF_NOISE_PSK, CONF_PASSWORD, CONF_MAC_ADDRESS}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: ESPHomeConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
     diag: dict[str, Any] = {}
 
     diag["config"] = config_entry.as_dict()
 
-    entry_data = DomainData.get(hass).get_entry_data(config_entry)
+    entry_data = config_entry.runtime_data
 
     if (storage_data := await entry_data.store.async_load()) is not None:
         diag["storage_data"] = storage_data

homeassistant/components/esphome/domain_data.py

@@ -3,16 +3,16 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Self, cast
+from functools import cache
+from typing import Self
 
 from bleak_esphome.backend.cache import ESPHomeBluetoothCache
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.json import JSONEncoder
 
 from .const import DOMAIN
-from .entry_data import ESPHomeStorage, RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, ESPHomeStorage, RuntimeEntryData
 
 STORAGE_VERSION = 1
 
@@ -21,30 +21,26 @@ STORAGE_VERSION = 1
 class DomainData:
     """Define a class that stores global esphome data in hass.data[DOMAIN]."""
 
-    _entry_datas: dict[str, RuntimeEntryData] = field(default_factory=dict)
     _stores: dict[str, ESPHomeStorage] = field(default_factory=dict)
     bluetooth_cache: ESPHomeBluetoothCache = field(
         default_factory=ESPHomeBluetoothCache
     )
 
-    def get_entry_data(self, entry: ConfigEntry) -> RuntimeEntryData:
+    def get_entry_data(self, entry: ESPHomeConfigEntry) -> RuntimeEntryData:
         """Return the runtime entry data associated with this config entry.
 
         Raises KeyError if the entry isn't loaded yet.
         """
-        return self._entry_datas[entry.entry_id]
+        return entry.runtime_data
 
-    def set_entry_data(self, entry: ConfigEntry, entry_data: RuntimeEntryData) -> None:
+    def set_entry_data(
+        self, entry: ESPHomeConfigEntry, entry_data: RuntimeEntryData
+    ) -> None:
         """Set the runtime entry data associated with this config entry."""
-        assert entry.entry_id not in self._entry_datas, "Entry data already set!"
-        self._entry_datas[entry.entry_id] = entry_data
-
-    def pop_entry_data(self, entry: ConfigEntry) -> RuntimeEntryData:
-        """Pop the runtime entry data instance associated with this config entry."""
-        return self._entry_datas.pop(entry.entry_id)
+        entry.runtime_data = entry_data
 
     def get_or_create_store(
-        self, hass: HomeAssistant, entry: ConfigEntry
+        self, hass: HomeAssistant, entry: ESPHomeConfigEntry
     ) -> ESPHomeStorage:
         """Get or create a Store instance for the given config entry."""
         return self._stores.setdefault(
@@ -55,10 +51,8 @@ class DomainData:
         )
 
     @classmethod
+    @cache
     def get(cls, hass: HomeAssistant) -> Self:
         """Get the global DomainData instance stored in hass.data."""
-        # Don't use setdefault - this is a hot code path
-        if DOMAIN in hass.data:
-            return cast(Self, hass.data[DOMAIN])
         ret = hass.data[DOMAIN] = cls()
         return ret

homeassistant/components/esphome/entity.py

@@ -16,7 +16,6 @@ from aioesphomeapi import (
 )
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
@@ -27,10 +26,8 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .domain_data import DomainData
-
 # Import config flow so that it's added to the registry
-from .entry_data import RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
 from .enum_mapper import EsphomeEnumMapper
 
 _R = TypeVar("_R")
@@ -85,7 +82,7 @@ def async_static_info_updated(
 
 async def platform_async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: ESPHomeConfigEntry,
     async_add_entities: AddEntitiesCallback,
     *,
     info_type: type[_InfoT],
@@ -97,7 +94,7 @@ async def platform_async_setup_entry(
     This method is in charge of receiving, distributing and storing
     info and state updates.
     """
-    entry_data: RuntimeEntryData = DomainData.get(hass).get_entry_data(entry)
+    entry_data = entry.runtime_data
     entry_data.info[info_type] = {}
     entry_data.state.setdefault(state_type, {})
     platform = entity_platform.async_get_current_platform()

homeassistant/components/esphome/entry_data.py

@@ -55,6 +55,9 @@ from homeassistant.helpers.storage import Store
 from .const import DOMAIN
 from .dashboard import async_get_dashboard
 
+type ESPHomeConfigEntry = ConfigEntry[RuntimeEntryData]
+
+
 INFO_TO_COMPONENT_TYPE: Final = {v: k for k, v in COMPONENT_TYPE_TO_INFO.items()}
 
 _SENTINEL = object()
@@ -248,7 +251,7 @@ class RuntimeEntryData:
     async def _ensure_platforms_loaded(
         self,
         hass: HomeAssistant,
-        entry: ConfigEntry,
+        entry: ESPHomeConfigEntry,
         platforms: set[Platform],
     ) -> None:
         async with self.platform_load_lock:
@@ -259,7 +262,7 @@ class RuntimeEntryData:
     async def async_update_static_infos(
         self,
         hass: HomeAssistant,
-        entry: ConfigEntry,
+        entry: ESPHomeConfigEntry,
         infos: list[EntityInfo],
         mac: str,
     ) -> None:
@@ -452,7 +455,7 @@ class RuntimeEntryData:
             await self.store.async_save(self._pending_storage())
 
     async def async_update_listener(
-        self, hass: HomeAssistant, entry: ConfigEntry
+        self, hass: HomeAssistant, entry: ESPHomeConfigEntry
     ) -> None:
         """Handle options update."""
         if self.original_options == entry.options:

homeassistant/components/esphome/event.py

@@ -5,16 +5,18 @@ from __future__ import annotations
 from aioesphomeapi import EntityInfo, Event, EventInfo
 
 from homeassistant.components.event import EventDeviceClass, EventEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
 
 from .entity import EsphomeEntity, platform_async_setup_entry
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome event based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/fan.py

@@ -13,7 +13,6 @@ from homeassistant.components.fan import (
     FanEntity,
     FanEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.percentage import (
@@ -29,13 +28,16 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 from .enum_mapper import EsphomeEnumMapper
 
 ORDERED_NAMED_FAN_SPEEDS = [FanSpeed.LOW, FanSpeed.MEDIUM, FanSpeed.HIGH]
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome fans based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/light.py

@@ -29,7 +29,6 @@ from homeassistant.components.light import (
     LightEntity,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
@@ -39,12 +38,15 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 FLASH_LENGTHS = {FLASH_SHORT: 2, FLASH_LONG: 10}
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome lights based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/lock.py

@@ -7,7 +7,6 @@ from typing import Any
 from aioesphomeapi import EntityInfo, LockCommand, LockEntityState, LockInfo, LockState
 
 from homeassistant.components.lock import LockEntity, LockEntityFeature
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_CODE
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -18,10 +17,13 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome switches based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/manager.py

@@ -28,7 +28,6 @@ import voluptuous as vol
 
 from homeassistant.components import tag, zeroconf
 from homeassistant.components.intent import async_register_timer_handler
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_DEVICE_ID,
     CONF_MODE,
@@ -73,7 +72,7 @@ from .dashboard import async_get_dashboard
 from .domain_data import DomainData
 
 # Import config flow so that it's added to the registry
-from .entry_data import RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
 from .voice_assistant import (
     VoiceAssistantAPIPipeline,
     VoiceAssistantPipeline,
@@ -159,7 +158,7 @@ class ESPHomeManager:
     def __init__(
         self,
         hass: HomeAssistant,
-        entry: ConfigEntry,
+        entry: ESPHomeConfigEntry,
         host: str,
         password: str | None,
         cli: APIClient,
@@ -639,7 +638,7 @@ class ESPHomeManager:
 
 @callback
 def _async_setup_device_registry(
-    hass: HomeAssistant, entry: ConfigEntry, entry_data: RuntimeEntryData
+    hass: HomeAssistant, entry: ESPHomeConfigEntry, entry_data: RuntimeEntryData
 ) -> str:
     """Set up device registry feature for a particular config entry."""
     device_info = entry_data.device_info
@@ -839,10 +838,11 @@ def _setup_services(
         _async_register_service(hass, entry_data, device_info, service)
 
 
-async def cleanup_instance(hass: HomeAssistant, entry: ConfigEntry) -> RuntimeEntryData:
+async def cleanup_instance(
+    hass: HomeAssistant, entry: ESPHomeConfigEntry
+) -> RuntimeEntryData:
     """Cleanup the esphome client if it exists."""
-    domain_data = DomainData.get(hass)
-    data = domain_data.pop_entry_data(entry)
+    data = entry.runtime_data
     data.async_on_disconnect()
     for cleanup_callback in data.cleanup_callbacks:
         cleanup_callback()

homeassistant/components/esphome/select.py

@@ -9,12 +9,10 @@ from homeassistant.components.assist_pipeline.select import (
     VadSensitivitySelect,
 )
 from homeassistant.components.select import SelectEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import DOMAIN
-from .domain_data import DomainData
 from .entity import (
     EsphomeAssistEntity,
     EsphomeEntity,
@@ -22,12 +20,12 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
-from .entry_data import RuntimeEntryData
+from .entry_data import ESPHomeConfigEntry, RuntimeEntryData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: ESPHomeConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up esphome selects based on a config entry."""
@@ -40,7 +38,7 @@ async def async_setup_entry(
         state_type=SelectState,
     )
 
-    entry_data = DomainData.get(hass).get_entry_data(entry)
+    entry_data = entry.runtime_data
     assert entry_data.device_info is not None
     if entry_data.device_info.voice_assistant_feature_flags_compat(
         entry_data.api_version

homeassistant/components/esphome/switch.py

@@ -7,7 +7,6 @@ from typing import Any
 from aioesphomeapi import EntityInfo, SwitchInfo, SwitchState
 
 from homeassistant.components.switch import SwitchDeviceClass, SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
@@ -18,10 +17,13 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome switches based on a config entry."""
     await platform_async_setup_entry(

homeassistant/components/esphome/valve.py

@@ -11,7 +11,6 @@ from homeassistant.components.valve import (
     ValveEntity,
     ValveEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.enum import try_parse_enum
@@ -22,10 +21,13 @@ from .entity import (
     esphome_state_property,
     platform_async_setup_entry,
 )
+from .entry_data import ESPHomeConfigEntry
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: ESPHomeConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up ESPHome valves based on a config entry."""
     await platform_async_setup_entry(