Migrate Wiz to config entry runtime data (#122091)

homeassistant/components/wiz/__init__.py

@@ -31,6 +31,8 @@ from .const import (
 from .discovery import async_discover_devices, async_trigger_discovery
 from .models import WizData
 
+type WizConfigEntry = ConfigEntry[WizData]
+
 _LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [
@@ -135,9 +137,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     await bulb.start_push(_async_push_update)
     bulb.set_discovery_callback(lambda bulb: async_trigger_discovery(hass, [bulb]))
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = WizData(
-        coordinator=coordinator, bulb=bulb, scenes=scenes
-    )
+    entry.runtime_data = WizData(coordinator=coordinator, bulb=bulb, scenes=scenes)
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     entry.async_on_unload(entry.add_update_listener(_async_update_listener))
@@ -147,6 +147,5 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload a config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        data: WizData = hass.data[DOMAIN].pop(entry.entry_id)
-        await data.bulb.async_close()
+        await entry.runtime_data.bulb.async_close()
     return unload_ok

homeassistant/components/wiz/binary_sensor.py

@@ -10,13 +10,13 @@ from homeassistant.components.binary_sensor import (
     BinarySensorDeviceClass,
     BinarySensorEntity,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
+from . import WizConfigEntry
 from .const import DOMAIN, SIGNAL_WIZ_PIR
 from .entity import WizEntity
 from .models import WizData
@@ -26,17 +26,16 @@ OCCUPANCY_UNIQUE_ID = "{}_occupancy"
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WizConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the WiZ binary sensor platform."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
-    mac = wiz_data.bulb.mac
+    mac = entry.runtime_data.bulb.mac
 
     if er.async_get(hass).async_get_entity_id(
         Platform.BINARY_SENSOR, DOMAIN, OCCUPANCY_UNIQUE_ID.format(mac)
     ):
-        async_add_entities([WizOccupancyEntity(wiz_data, entry.title)])
+        async_add_entities([WizOccupancyEntity(entry.runtime_data, entry.title)])
         return
 
     cancel_dispatcher: Callable[[], None] | None = None
@@ -47,7 +46,7 @@ async def async_setup_entry(
         assert cancel_dispatcher is not None
         cancel_dispatcher()
         cancel_dispatcher = None
-        async_add_entities([WizOccupancyEntity(wiz_data, entry.title)])
+        async_add_entities([WizOccupancyEntity(entry.runtime_data, entry.title)])
 
     cancel_dispatcher = async_dispatcher_connect(
         hass, SIGNAL_WIZ_PIR.format(mac), _async_add_occupancy_sensor

homeassistant/components/wiz/diagnostics.py

@@ -5,24 +5,21 @@ from __future__ import annotations
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .models import WizData
+from . import WizConfigEntry
 
 TO_REDACT = {"roomId", "homeId"}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: WizConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
     return {
         "entry": {
             "title": entry.title,
             "data": dict(entry.data),
         },
-        "data": async_redact_data(wiz_data.bulb.diagnostics, TO_REDACT),
+        "data": async_redact_data(entry.runtime_data.bulb.diagnostics, TO_REDACT),
     }

homeassistant/components/wiz/light.py

@@ -19,7 +19,6 @@ from homeassistant.components.light import (
     LightEntityFeature,
     filter_supported_color_modes,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import (
@@ -27,7 +26,7 @@ from homeassistant.util.color import (
     color_temperature_mired_to_kelvin,
 )
 
-from .const import DOMAIN
+from . import WizConfigEntry
 from .entity import WizToggleEntity
 from .models import WizData
 
@@ -61,13 +60,12 @@ def _async_pilot_builder(**kwargs: Any) -> PilotBuilder:
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WizConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the WiZ Platform from config_flow."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
-    if wiz_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET:
-        async_add_entities([WizBulbEntity(wiz_data, entry.title)])
+    if entry.runtime_data.bulb.bulbtype.bulb_type != BulbClass.SOCKET:
+        async_add_entities([WizBulbEntity(entry.runtime_data, entry.title)])
 
 
 class WizBulbEntity(WizToggleEntity, LightEntity):

homeassistant/components/wiz/number.py

@@ -13,12 +13,11 @@ from homeassistant.components.number import (
     NumberEntityDescription,
     NumberMode,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import WizConfigEntry
 from .entity import WizEntity
 from .models import WizData
 
@@ -68,15 +67,16 @@ NUMBERS: tuple[WizNumberEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WizConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the wiz speed number."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
     async_add_entities(
-        WizSpeedNumber(wiz_data, entry.title, description)
+        WizSpeedNumber(entry.runtime_data, entry.title, description)
         for description in NUMBERS
-        if getattr(wiz_data.bulb.bulbtype.features, description.required_feature)
+        if getattr(
+            entry.runtime_data.bulb.bulbtype.features, description.required_feature
+        )
     )
 
 

homeassistant/components/wiz/sensor.py

@@ -8,7 +8,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     SIGNAL_STRENGTH_DECIBELS_MILLIWATT,
     EntityCategory,
@@ -17,7 +16,7 @@ from homeassistant.const import (
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import WizConfigEntry
 from .entity import WizEntity
 from .models import WizData
 
@@ -45,18 +44,18 @@ POWER_SENSORS: tuple[SensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WizConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the wiz sensor."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
     entities = [
-        WizSensor(wiz_data, entry.title, description) for description in SENSORS
+        WizSensor(entry.runtime_data, entry.title, description)
+        for description in SENSORS
     ]
-    if wiz_data.coordinator.data is not None:
+    if entry.runtime_data.coordinator.data is not None:
         entities.extend(
             [
-                WizPowerSensor(wiz_data, entry.title, description)
+                WizPowerSensor(entry.runtime_data, entry.title, description)
                 for description in POWER_SENSORS
             ]
         )

homeassistant/components/wiz/switch.py

@@ -8,24 +8,22 @@ from pywizlight import PilotBuilder
 from pywizlight.bulblibrary import BulbClass
 
 from homeassistant.components.switch import SwitchEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import WizConfigEntry
 from .entity import WizToggleEntity
 from .models import WizData
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: WizConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the WiZ switch platform."""
-    wiz_data: WizData = hass.data[DOMAIN][entry.entry_id]
-    if wiz_data.bulb.bulbtype.bulb_type == BulbClass.SOCKET:
-        async_add_entities([WizSocketEntity(wiz_data, entry.title)])
+    if entry.runtime_data.bulb.bulbtype.bulb_type == BulbClass.SOCKET:
+        async_add_entities([WizSocketEntity(entry.runtime_data, entry.title)])
 
 
 class WizSocketEntity(WizToggleEntity, SwitchEntity):