Migrate lifx to use runtime_data and HassKey (#147348)

homeassistant/components/lifx/__init__.py

@@ -13,7 +13,6 @@ from aiolifx.connection import LIFXConnection
 import voluptuous as vol
 
 from homeassistant.components.light import DOMAIN as LIGHT_DOMAIN
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     CONF_HOST,
     CONF_PORT,
@@ -27,7 +26,7 @@ from homeassistant.helpers.event import async_call_later, async_track_time_inter
 from homeassistant.helpers.typing import ConfigType
 
 from .const import _LOGGER, DATA_LIFX_MANAGER, DOMAIN, TARGET_ANY
-from .coordinator import LIFXUpdateCoordinator
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .discovery import async_discover_devices, async_trigger_discovery
 from .manager import LIFXManager
 from .migration import async_migrate_entities_devices, async_migrate_legacy_entries
@@ -73,7 +72,7 @@ DISCOVERY_COOLDOWN = 5
 
 async def async_legacy_migration(
     hass: HomeAssistant,
-    legacy_entry: ConfigEntry,
+    legacy_entry: LIFXConfigEntry,
     discovered_devices: Iterable[Light],
 ) -> bool:
     """Migrate config entries."""
@@ -157,7 +156,6 @@ class LIFXDiscoveryManager:
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up the LIFX component."""
-    hass.data[DOMAIN] = {}
     migrating = bool(async_get_legacy_entry(hass))
     discovery_manager = LIFXDiscoveryManager(hass, migrating)
 
@@ -187,7 +185,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: LIFXConfigEntry) -> bool:
     """Set up LIFX from a config entry."""
     if async_entry_is_legacy(entry):
         return True
@@ -198,10 +196,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         async_migrate_entities_devices(hass, legacy_entry.entry_id, entry)
 
     assert entry.unique_id is not None
-    domain_data = hass.data[DOMAIN]
-    if DATA_LIFX_MANAGER not in domain_data:
+    if DATA_LIFX_MANAGER not in hass.data:
         manager = LIFXManager(hass)
-        domain_data[DATA_LIFX_MANAGER] = manager
+        hass.data[DATA_LIFX_MANAGER] = manager
         manager.async_setup()
 
     host = entry.data[CONF_HOST]
@@ -229,21 +226,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
         raise ConfigEntryNotReady(
             f"Unexpected device found at {host}; expected {entry.unique_id}, found {serial}"
         )
-    domain_data[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: LIFXConfigEntry) -> bool:
     """Unload a config entry."""
     if async_entry_is_legacy(entry):
         return True
-    domain_data = hass.data[DOMAIN]
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        coordinator: LIFXUpdateCoordinator = domain_data.pop(entry.entry_id)
-        coordinator.connection.async_stop()
+        entry.runtime_data.connection.async_stop()
     # Only the DATA_LIFX_MANAGER left, remove it.
-    if len(domain_data) == 1:
-        manager: LIFXManager = domain_data.pop(DATA_LIFX_MANAGER)
+    if len(hass.config_entries.async_loaded_entries(DOMAIN)) == 0:
+        manager = hass.data.pop(DATA_LIFX_MANAGER)
         manager.async_unload()
     return unload_ok

homeassistant/components/lifx/binary_sensor.py

@@ -7,13 +7,12 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, HEV_CYCLE_STATE
-from .coordinator import LIFXUpdateCoordinator
+from .const import HEV_CYCLE_STATE
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .entity import LIFXEntity
 from .util import lifx_features
 
@@ -27,11 +26,11 @@ HEV_CYCLE_STATE_SENSOR = BinarySensorEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: LIFXConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up LIFX from a config entry."""
-    coordinator: LIFXUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     if lifx_features(coordinator.device)["hev"]:
         async_add_entities(

homeassistant/components/lifx/button.py

@@ -7,13 +7,12 @@ from homeassistant.components.button import (
     ButtonEntity,
     ButtonEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN, IDENTIFY, RESTART
-from .coordinator import LIFXUpdateCoordinator
+from .const import IDENTIFY, RESTART
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .entity import LIFXEntity
 
 RESTART_BUTTON_DESCRIPTION = ButtonEntityDescription(
@@ -31,12 +30,11 @@ IDENTIFY_BUTTON_DESCRIPTION = ButtonEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: LIFXConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up LIFX from a config entry."""
-    domain_data = hass.data[DOMAIN]
-    coordinator: LIFXUpdateCoordinator = domain_data[entry.entry_id]
+    coordinator = entry.runtime_data
     async_add_entities(
         [LIFXRestartButton(coordinator), LIFXIdentifyButton(coordinator)]
     )

homeassistant/components/lifx/const.py

@@ -1,8 +1,17 @@
 """Const for LIFX."""
 
+from __future__ import annotations
+
 import logging
+from typing import TYPE_CHECKING
+
+from homeassistant.util.hass_dict import HassKey
+
+if TYPE_CHECKING:
+    from .manager import LIFXManager
 
 DOMAIN = "lifx"
+DATA_LIFX_MANAGER: HassKey[LIFXManager] = HassKey(DOMAIN)
 
 TARGET_ANY = "00:00:00:00:00:00"
 
@@ -59,7 +68,6 @@ INFRARED_BRIGHTNESS_VALUES_MAP = {
     32767: "50%",
     65535: "100%",
 }
-DATA_LIFX_MANAGER = "lifx_manager"
 
 LIFX_CEILING_PRODUCT_IDS = {176, 177, 201, 202}
 

homeassistant/components/lifx/coordinator.py

@@ -65,6 +65,8 @@ ZONES_PER_COLOR_UPDATE_REQUEST = 8
 
 RSSI_DBM_FW = AwesomeVersion("2.77")
 
+type LIFXConfigEntry = ConfigEntry[LIFXUpdateCoordinator]
+
 
 class FirmwareEffect(IntEnum):
     """Enumeration of LIFX firmware effects."""
@@ -87,12 +89,12 @@ class SkyType(IntEnum):
 class LIFXUpdateCoordinator(DataUpdateCoordinator[None]):
     """DataUpdateCoordinator to gather data for a specific lifx device."""
 
-    config_entry: ConfigEntry
+    config_entry: LIFXConfigEntry
 
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: LIFXConfigEntry,
         connection: LIFXConnection,
     ) -> None:
         """Initialize DataUpdateCoordinator."""

homeassistant/components/lifx/diagnostics.py

@@ -5,21 +5,20 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, CONF_IP_ADDRESS, CONF_MAC
 from homeassistant.core import HomeAssistant
 
-from .const import CONF_LABEL, DOMAIN
-from .coordinator import LIFXUpdateCoordinator
+from .const import CONF_LABEL
+from .coordinator import LIFXConfigEntry
 
 TO_REDACT = [CONF_LABEL, CONF_HOST, CONF_IP_ADDRESS, CONF_MAC]
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: LIFXConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a LIFX config entry."""
-    coordinator: LIFXUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     return {
         "entry": {
             "title": entry.title,

homeassistant/components/lifx/light.py

@@ -17,7 +17,6 @@ from homeassistant.components.light import (
     LightEntity,
     LightEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_ENTITY_ID, Platform
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -37,7 +36,7 @@ from .const import (
     INFRARED_BRIGHTNESS,
     LIFX_CEILING_PRODUCT_IDS,
 )
-from .coordinator import FirmwareEffect, LIFXUpdateCoordinator
+from .coordinator import FirmwareEffect, LIFXConfigEntry, LIFXUpdateCoordinator
 from .entity import LIFXEntity
 from .manager import (
     SERVICE_EFFECT_COLORLOOP,
@@ -78,13 +77,12 @@ HSBK_KELVIN = 3
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: LIFXConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up LIFX from a config entry."""
-    domain_data = hass.data[DOMAIN]
-    coordinator: LIFXUpdateCoordinator = domain_data[entry.entry_id]
-    manager: LIFXManager = domain_data[DATA_LIFX_MANAGER]
+    coordinator = entry.runtime_data
+    manager = hass.data[DATA_LIFX_MANAGER]
     device = coordinator.device
     platform = entity_platform.async_get_current_platform()
     platform.async_register_entity_service(
@@ -123,7 +121,7 @@ class LIFXLight(LIFXEntity, LightEntity):
         self,
         coordinator: LIFXUpdateCoordinator,
         manager: LIFXManager,
-        entry: ConfigEntry,
+        entry: LIFXConfigEntry,
     ) -> None:
         """Initialize the light."""
         super().__init__(coordinator)

homeassistant/components/lifx/manager.py

@@ -30,8 +30,8 @@ from homeassistant.core import HomeAssistant, ServiceCall, callback
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.service import async_extract_referenced_entity_ids
 
-from .const import _ATTR_COLOR_TEMP, ATTR_THEME, DATA_LIFX_MANAGER, DOMAIN
-from .coordinator import LIFXUpdateCoordinator
+from .const import _ATTR_COLOR_TEMP, ATTR_THEME, DOMAIN
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .util import convert_8_to_16, find_hsbk
 
 if TYPE_CHECKING:
@@ -494,13 +494,11 @@ class LIFXManager:
         coordinators: list[LIFXUpdateCoordinator] = []
         bulbs: list[Light] = []
 
-        for entry_id, coordinator in self.hass.data[DOMAIN].items():
-            if (
-                entry_id != DATA_LIFX_MANAGER
-                and self.entry_id_to_entity_id[entry_id] in entity_ids
-            ):
-                coordinators.append(coordinator)
-                bulbs.append(coordinator.device)
+        entry: LIFXConfigEntry
+        for entry in self.hass.config_entries.async_loaded_entries(DOMAIN):
+            if self.entry_id_to_entity_id[entry.entry_id] in entity_ids:
+                coordinators.append(entry.runtime_data)
+                bulbs.append(entry.runtime_data.device)
 
         if start_effect_func := self._effect_dispatch.get(service):
             await start_effect_func(self, bulbs, coordinators, **kwargs)

homeassistant/components/lifx/migration.py

@@ -2,11 +2,11 @@
 
 from __future__ import annotations
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 
 from .const import _LOGGER, DOMAIN
+from .coordinator import LIFXConfigEntry
 from .discovery import async_init_discovery_flow
 
 
@@ -15,7 +15,7 @@ def async_migrate_legacy_entries(
     hass: HomeAssistant,
     discovered_hosts_by_serial: dict[str, str],
     existing_serials: set[str],
-    legacy_entry: ConfigEntry,
+    legacy_entry: LIFXConfigEntry,
 ) -> int:
     """Migrate the legacy config entries to have an entry per device."""
     _LOGGER.debug(
@@ -45,7 +45,7 @@ def async_migrate_legacy_entries(
 
 @callback
 def async_migrate_entities_devices(
-    hass: HomeAssistant, legacy_entry_id: str, new_entry: ConfigEntry
+    hass: HomeAssistant, legacy_entry_id: str, new_entry: LIFXConfigEntry
 ) -> None:
     """Move entities and devices to the new config entry."""
     migrated_devices = []

homeassistant/components/lifx/select.py

@@ -5,18 +5,12 @@ from __future__ import annotations
 from aiolifx_themes.themes import ThemeLibrary
 
 from homeassistant.components.select import SelectEntity, SelectEntityDescription
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import (
-    ATTR_THEME,
-    DOMAIN,
-    INFRARED_BRIGHTNESS,
-    INFRARED_BRIGHTNESS_VALUES_MAP,
-)
-from .coordinator import LIFXUpdateCoordinator
+from .const import ATTR_THEME, INFRARED_BRIGHTNESS, INFRARED_BRIGHTNESS_VALUES_MAP
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .entity import LIFXEntity
 from .util import lifx_features
 
@@ -39,11 +33,11 @@ THEME_ENTITY = SelectEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: LIFXConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up LIFX from a config entry."""
-    coordinator: LIFXUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
 
     entities: list[LIFXEntity] = []
 

homeassistant/components/lifx/sensor.py

@@ -10,13 +10,12 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import ATTR_RSSI, DOMAIN
-from .coordinator import LIFXUpdateCoordinator
+from .const import ATTR_RSSI
+from .coordinator import LIFXConfigEntry, LIFXUpdateCoordinator
 from .entity import LIFXEntity
 
 SCAN_INTERVAL = timedelta(seconds=30)
@@ -33,11 +32,11 @@ RSSI_SENSOR = SensorEntityDescription(
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: LIFXConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up LIFX sensor from config entry."""
-    coordinator: LIFXUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    coordinator = entry.runtime_data
     async_add_entities([LIFXRssiSensor(coordinator, RSSI_SENSOR)])
 
 

homeassistant/components/lifx/util.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import Callable
 from functools import partial
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from aiolifx import products
 from aiolifx.aiolifx import Light
@@ -21,7 +21,6 @@ from homeassistant.components.light import (
     ATTR_RGB_COLOR,
     ATTR_XY_COLOR,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
 from homeassistant.util import color as color_util
@@ -35,17 +34,20 @@ from .const import (
     OVERALL_TIMEOUT,
 )
 
+if TYPE_CHECKING:
+    from .coordinator import LIFXConfigEntry
+
 FIX_MAC_FW = AwesomeVersion("3.70")
 
 
 @callback
-def async_entry_is_legacy(entry: ConfigEntry) -> bool:
+def async_entry_is_legacy(entry: LIFXConfigEntry) -> bool:
     """Check if a config entry is the legacy shared one."""
     return entry.unique_id is None or entry.unique_id == DOMAIN
 
 
 @callback
-def async_get_legacy_entry(hass: HomeAssistant) -> ConfigEntry | None:
+def async_get_legacy_entry(hass: HomeAssistant) -> LIFXConfigEntry | None:
     """Get the legacy config entry."""
     for entry in hass.config_entries.async_entries(DOMAIN):
         if async_entry_is_legacy(entry):