Migrate HomeWizard to config entry runtime data (#122088)

homeassistant/components/homewizard/__init__.py

@@ -5,12 +5,14 @@ from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
 from .const import DOMAIN, PLATFORMS
-from .coordinator import HWEnergyDeviceUpdateCoordinator as Coordinator
+from .coordinator import HWEnergyDeviceUpdateCoordinator
+
+type HomeWizardConfigEntry = ConfigEntry[HWEnergyDeviceUpdateCoordinator]
 
 
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Homewizard from a config entry."""
-    coordinator = Coordinator(hass)
+    coordinator = HWEnergyDeviceUpdateCoordinator(hass)
     try:
         await coordinator.async_config_entry_first_refresh()
 
@@ -22,7 +24,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
         raise
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
+    entry.runtime_data = coordinator
 
     # Abort reauth config flow if active
     for progress_flow in hass.config_entries.flow.async_progress_by_handler(DOMAIN):
@@ -41,6 +43,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):
-        coordinator = hass.data[DOMAIN].pop(entry.entry_id)
-        await coordinator.api.close()
+        entry.runtime_data.api.close()
     return unload_ok

homeassistant/components/homewizard/button.py

@@ -1,24 +1,24 @@
 """Support for HomeWizard buttons."""
 
 from homeassistant.components.button import ButtonDeviceClass, ButtonEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN
+from . import HomeWizardConfigEntry
 from .coordinator import HWEnergyDeviceUpdateCoordinator
 from .entity import HomeWizardEntity
 from .helpers import homewizard_exception_handler
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HomeWizardConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up the Identify button."""
-    coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    if coordinator.supports_identify():
-        async_add_entities([HomeWizardIdentifyButton(coordinator)])
+    if entry.runtime_data.supports_identify():
+        async_add_entities([HomeWizardIdentifyButton(entry.runtime_data)])
 
 
 class HomeWizardIdentifyButton(HomeWizardEntity, ButtonEntity):

homeassistant/components/homewizard/diagnostics.py

@@ -6,12 +6,10 @@ from dataclasses import asdict
 from typing import Any
 
 from homeassistant.components.diagnostics import async_redact_data
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_IP_ADDRESS
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
-from .coordinator import HWEnergyDeviceUpdateCoordinator
+from . import HomeWizardConfigEntry
 
 TO_REDACT = {
     CONF_IP_ADDRESS,
@@ -24,25 +22,25 @@ TO_REDACT = {
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, entry: ConfigEntry
+    hass: HomeAssistant, entry: HomeWizardConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    data = entry.runtime_data.data
 
     state: dict[str, Any] | None = None
-    if coordinator.data.state:
-        state = asdict(coordinator.data.state)
+    if data.state:
+        state = asdict(data.state)
 
     system: dict[str, Any] | None = None
-    if coordinator.data.system:
-        system = asdict(coordinator.data.system)
+    if data.system:
+        system = asdict(data.system)
 
     return async_redact_data(
         {
             "entry": async_redact_data(entry.data, TO_REDACT),
             "data": {
-                "device": asdict(coordinator.data.device),
-                "data": asdict(coordinator.data.data),
+                "device": asdict(data.device),
+                "data": asdict(data.data),
                 "state": state,
                 "system": system,
             },

homeassistant/components/homewizard/number.py

@@ -3,13 +3,12 @@
 from __future__ import annotations
 
 from homeassistant.components.number import NumberEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import PERCENTAGE, EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util.color import brightness_to_value, value_to_brightness
 
-from .const import DOMAIN
+from . import HomeWizardConfigEntry
 from .coordinator import HWEnergyDeviceUpdateCoordinator
 from .entity import HomeWizardEntity
 from .helpers import homewizard_exception_handler
@@ -17,13 +16,12 @@ from .helpers import homewizard_exception_handler
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: HomeWizardConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up numbers for device."""
-    coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-    if coordinator.supports_state():
-        async_add_entities([HWEnergyNumberEntity(coordinator)])
+    if entry.runtime_data.supports_state():
+        async_add_entities([HWEnergyNumberEntity(entry.runtime_data)])
 
 
 class HWEnergyNumberEntity(HomeWizardEntity, NumberEntity):

homeassistant/components/homewizard/sensor.py

@@ -15,7 +15,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_VIA_DEVICE,
     PERCENTAGE,
@@ -36,6 +35,7 @@ from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 
+from . import HomeWizardConfigEntry
 from .const import DOMAIN
 from .coordinator import HWEnergyDeviceUpdateCoordinator
 from .entity import HomeWizardEntity
@@ -619,24 +619,25 @@ EXTERNAL_SENSORS = {
 
 
 async def async_setup_entry(
-    hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
+    hass: HomeAssistant,
+    entry: HomeWizardConfigEntry,
+    async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Initialize sensors."""
-    coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
 
     # Migrate original gas meter sensor to ExternalDevice
     # This is sensor that was directly linked to the P1 Meter
     # Migration can be removed after 2024.8.0
     ent_reg = er.async_get(hass)
-
+    data = entry.runtime_data.data.data
     if (
         entity_id := ent_reg.async_get_entity_id(
             Platform.SENSOR, DOMAIN, f"{entry.unique_id}_total_gas_m3"
         )
-    ) and coordinator.data.data.gas_unique_id is not None:
+    ) and data.gas_unique_id is not None:
         ent_reg.async_update_entity(
             entity_id,
-            new_unique_id=f"{DOMAIN}_gas_meter_{coordinator.data.data.gas_unique_id}",
+            new_unique_id=f"{DOMAIN}_gas_meter_{data.gas_unique_id}",
         )
 
     # Remove old gas_unique_id sensor
@@ -647,14 +648,14 @@ async def async_setup_entry(
 
     # Initialize default sensors
     entities: list = [
-        HomeWizardSensorEntity(coordinator, description)
+        HomeWizardSensorEntity(entry.runtime_data, description)
         for description in SENSORS
-        if description.has_fn(coordinator.data.data)
+        if description.has_fn(data)
     ]
 
     # Initialize external devices
-    if coordinator.data.data.external_devices is not None:
-        for unique_id, device in coordinator.data.data.external_devices.items():
+    if data.external_devices is not None:
+        for unique_id, device in data.external_devices.items():
             if description := EXTERNAL_SENSORS.get(device.meter_type):
                 # Migrate external devices to new unique_id
                 # This is to ensure that devices with same id but different type are unique
@@ -669,7 +670,9 @@ async def async_setup_entry(
 
                 # Add external device
                 entities.append(
-                    HomeWizardExternalSensorEntity(coordinator, description, unique_id)
+                    HomeWizardExternalSensorEntity(
+                        entry.runtime_data, description, unique_id
+                    )
                 )
 
     async_add_entities(entities)

homeassistant/components/homewizard/switch.py

@@ -13,12 +13,12 @@ from homeassistant.components.switch import (
     SwitchEntity,
     SwitchEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EntityCategory
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import DOMAIN, DeviceResponseEntry
+from . import HomeWizardConfigEntry
+from .const import DeviceResponseEntry
 from .coordinator import HWEnergyDeviceUpdateCoordinator
 from .entity import HomeWizardEntity
 from .helpers import homewizard_exception_handler
@@ -67,16 +67,14 @@ SWITCHES = [
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    entry: ConfigEntry,
+    entry: HomeWizardConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up switches."""
-    coordinator: HWEnergyDeviceUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
-
     async_add_entities(
-        HomeWizardSwitchEntity(coordinator, description)
+        HomeWizardSwitchEntity(entry.runtime_data, description)
         for description in SWITCHES
-        if description.create_fn(coordinator)
+        if description.create_fn(entry.runtime_data)
     )