Use runtime_data in hydrawise (#144950)

homeassistant/components/hydrawise/__init__.py

@@ -2,13 +2,13 @@
 
 from pydrawise import auth, hybrid
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_PASSWORD, CONF_USERNAME, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryAuthFailed
 
-from .const import APP_ID, DOMAIN
+from .const import APP_ID
 from .coordinator import (
+    HydrawiseConfigEntry,
     HydrawiseMainDataUpdateCoordinator,
     HydrawiseUpdateCoordinators,
     HydrawiseWaterUseDataUpdateCoordinator,
@@ -24,7 +24,9 @@ PLATFORMS: list[Platform] = [
 _REQUIRED_AUTH_KEYS = (CONF_USERNAME, CONF_PASSWORD, CONF_API_KEY)
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(
+    hass: HomeAssistant, config_entry: HydrawiseConfigEntry
+) -> bool:
     """Set up Hydrawise from a config entry."""
     if any(k not in config_entry.data for k in _REQUIRED_AUTH_KEYS):
         # If we are missing any required authentication keys, trigger a reauth flow.
@@ -45,18 +47,14 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
         hass, config_entry, hydrawise, main_coordinator
     )
     await water_use_coordinator.async_config_entry_first_refresh()
-    hass.data.setdefault(DOMAIN, {})[config_entry.entry_id] = (
-        HydrawiseUpdateCoordinators(
-            main=main_coordinator,
-            water_use=water_use_coordinator,
-        )
+    config_entry.runtime_data = HydrawiseUpdateCoordinators(
+        main=main_coordinator,
+        water_use=water_use_coordinator,
     )
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: HydrawiseConfigEntry) -> bool:
     """Unload a config entry."""
-    if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        hass.data[DOMAIN].pop(entry.entry_id)
-    return unload_ok
+    return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/hydrawise/binary_sensor.py

@@ -14,14 +14,13 @@ from homeassistant.components.binary_sensor import (
     BinarySensorEntity,
     BinarySensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import config_validation as cv, entity_platform
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.typing import VolDictType
 
-from .const import DOMAIN, SERVICE_RESUME, SERVICE_START_WATERING, SERVICE_SUSPEND
-from .coordinator import HydrawiseUpdateCoordinators
+from .const import SERVICE_RESUME, SERVICE_START_WATERING, SERVICE_SUSPEND
+from .coordinator import HydrawiseConfigEntry
 from .entity import HydrawiseEntity
 
 
@@ -77,11 +76,11 @@ SCHEMA_SUSPEND: VolDictType = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HydrawiseConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Hydrawise binary_sensor platform."""
-    coordinators: HydrawiseUpdateCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
     entities: list[HydrawiseBinarySensor] = []
     for controller in coordinators.main.data.controllers.values():
         entities.extend(

homeassistant/components/hydrawise/coordinator.py

@@ -14,6 +14,8 @@ from homeassistant.util.dt import now
 
 from .const import DOMAIN, LOGGER, MAIN_SCAN_INTERVAL, WATER_USE_SCAN_INTERVAL
 
+type HydrawiseConfigEntry = ConfigEntry[HydrawiseUpdateCoordinators]
+
 
 @dataclass
 class HydrawiseData:
@@ -40,7 +42,7 @@ class HydrawiseDataUpdateCoordinator(DataUpdateCoordinator[HydrawiseData]):
     """Base class for Hydrawise Data Update Coordinators."""
 
     api: HydrawiseBase
-    config_entry: ConfigEntry
+    config_entry: HydrawiseConfigEntry
 
 
 class HydrawiseMainDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
@@ -52,7 +54,10 @@ class HydrawiseMainDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
     """
 
     def __init__(
-        self, hass: HomeAssistant, config_entry: ConfigEntry, api: HydrawiseBase
+        self,
+        hass: HomeAssistant,
+        config_entry: HydrawiseConfigEntry,
+        api: HydrawiseBase,
     ) -> None:
         """Initialize HydrawiseDataUpdateCoordinator."""
         super().__init__(
@@ -92,7 +97,7 @@ class HydrawiseWaterUseDataUpdateCoordinator(HydrawiseDataUpdateCoordinator):
     def __init__(
         self,
         hass: HomeAssistant,
-        config_entry: ConfigEntry,
+        config_entry: HydrawiseConfigEntry,
         api: HydrawiseBase,
         main_coordinator: HydrawiseMainDataUpdateCoordinator,
     ) -> None:

homeassistant/components/hydrawise/sensor.py

@@ -14,14 +14,12 @@ from homeassistant.components.sensor import (
     SensorEntity,
     SensorEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import UnitOfTime, UnitOfVolume
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN
-from .coordinator import HydrawiseUpdateCoordinators
+from .coordinator import HydrawiseConfigEntry
 from .entity import HydrawiseEntity
 
 
@@ -130,11 +128,11 @@ FLOW_MEASUREMENT_KEYS = [x.key for x in FLOW_CONTROLLER_SENSORS]
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HydrawiseConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Hydrawise sensor platform."""
-    coordinators: HydrawiseUpdateCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
     entities: list[HydrawiseSensor] = []
     for controller in coordinators.main.data.controllers.values():
         entities.extend(

homeassistant/components/hydrawise/switch.py

@@ -14,13 +14,12 @@ from homeassistant.components.switch import (
     SwitchEntity,
     SwitchEntityDescription,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.util import dt as dt_util
 
-from .const import DEFAULT_WATERING_TIME, DOMAIN
-from .coordinator import HydrawiseUpdateCoordinators
+from .const import DEFAULT_WATERING_TIME
+from .coordinator import HydrawiseConfigEntry
 from .entity import HydrawiseEntity
 
 
@@ -62,11 +61,11 @@ SWITCH_KEYS: list[str] = [desc.key for desc in SWITCH_TYPES]
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HydrawiseConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Hydrawise switch platform."""
-    coordinators: HydrawiseUpdateCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
     async_add_entities(
         HydrawiseSwitch(coordinators.main, description, controller, zone_id=zone.id)
         for controller in coordinators.main.data.controllers.values()

homeassistant/components/hydrawise/valve.py

@@ -12,12 +12,10 @@ from homeassistant.components.valve import (
     ValveEntityDescription,
     ValveEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 
-from .const import DOMAIN
-from .coordinator import HydrawiseUpdateCoordinators
+from .coordinator import HydrawiseConfigEntry
 from .entity import HydrawiseEntity
 
 VALVE_TYPES: tuple[ValveEntityDescription, ...] = (
@@ -30,11 +28,11 @@ VALVE_TYPES: tuple[ValveEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: HydrawiseConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up the Hydrawise valve platform."""
-    coordinators: HydrawiseUpdateCoordinators = hass.data[DOMAIN][config_entry.entry_id]
+    coordinators = config_entry.runtime_data
     async_add_entities(
         HydrawiseValve(coordinators.main, description, controller, zone_id=zone.id)
         for controller in coordinators.main.data.controllers.values()