Use runtime_data in environment_canada (#136805)

homeassistant/components/environment_canada/__init__.py

@@ -5,14 +5,13 @@ import logging
 
 from env_canada import ECAirQuality, ECRadar, ECWeather
 
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_LANGUAGE, CONF_LATITUDE, CONF_LONGITUDE, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
 
 from .const import CONF_STATION, DOMAIN
-from .coordinator import ECDataUpdateCoordinator
+from .coordinator import ECConfigEntry, ECDataUpdateCoordinator, ECRuntimeData
 
 DEFAULT_RADAR_UPDATE_INTERVAL = timedelta(minutes=5)
 DEFAULT_WEATHER_UPDATE_INTERVAL = timedelta(minutes=5)
@@ -22,14 +21,13 @@ PLATFORMS = [Platform.CAMERA, Platform.SENSOR, Platform.WEATHER]
 _LOGGER = logging.getLogger(__name__)
 
 
-async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, config_entry: ECConfigEntry) -> bool:
     """Set up EC as config entry."""
     lat = config_entry.data.get(CONF_LATITUDE)
     lon = config_entry.data.get(CONF_LONGITUDE)
     station = config_entry.data.get(CONF_STATION)
     lang = config_entry.data.get(CONF_LANGUAGE, "English")
 
-    coordinators = {}
     errors = 0
 
     weather_data = ECWeather(
@@ -37,31 +35,31 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
         coordinates=(lat, lon),
         language=lang.lower(),
     )
-    coordinators["weather_coordinator"] = ECDataUpdateCoordinator(
+    weather_coordinator = ECDataUpdateCoordinator(
-        hass, weather_data, "weather", DEFAULT_WEATHER_UPDATE_INTERVAL
+        hass, config_entry, weather_data, "weather", DEFAULT_WEATHER_UPDATE_INTERVAL
     )
     try:
-        await coordinators["weather_coordinator"].async_config_entry_first_refresh()
+        await weather_coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady:
         errors = errors + 1
         _LOGGER.warning("Unable to retrieve Environment Canada weather")
 
     radar_data = ECRadar(coordinates=(lat, lon))
-    coordinators["radar_coordinator"] = ECDataUpdateCoordinator(
+    radar_coordinator = ECDataUpdateCoordinator(
-        hass, radar_data, "radar", DEFAULT_RADAR_UPDATE_INTERVAL
+        hass, config_entry, radar_data, "radar", DEFAULT_RADAR_UPDATE_INTERVAL
     )
     try:
-        await coordinators["radar_coordinator"].async_config_entry_first_refresh()
+        await radar_coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady:
         errors = errors + 1
         _LOGGER.warning("Unable to retrieve Environment Canada radar")
 
     aqhi_data = ECAirQuality(coordinates=(lat, lon))
-    coordinators["aqhi_coordinator"] = ECDataUpdateCoordinator(
+    aqhi_coordinator = ECDataUpdateCoordinator(
-        hass, aqhi_data, "AQHI", DEFAULT_WEATHER_UPDATE_INTERVAL
+        hass, config_entry, aqhi_data, "AQHI", DEFAULT_WEATHER_UPDATE_INTERVAL
     )
     try:
-        await coordinators["aqhi_coordinator"].async_config_entry_first_refresh()
+        await aqhi_coordinator.async_config_entry_first_refresh()
     except ConfigEntryNotReady:
         errors = errors + 1
         _LOGGER.warning("Unable to retrieve Environment Canada AQHI")
@@ -69,26 +67,23 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     if errors == 3:
         raise ConfigEntryNotReady
 
-    hass.data.setdefault(DOMAIN, {})
+    config_entry.runtime_data = ECRuntimeData(
-    hass.data[DOMAIN][config_entry.entry_id] = coordinators
+        aqhi_coordinator=aqhi_coordinator,
+        radar_coordinator=radar_coordinator,
+        weather_coordinator=weather_coordinator,
+    )
 
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, config_entry: ECConfigEntry) -> bool:
     """Unload a config entry."""
-    unload_ok = await hass.config_entries.async_unload_platforms(
+    return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS)
-        config_entry, PLATFORMS
-    )
-
-    hass.data[DOMAIN].pop(config_entry.entry_id)
-
-    return unload_ok
 
 
-def device_info(config_entry: ConfigEntry) -> DeviceInfo:
+def device_info(config_entry: ECConfigEntry) -> DeviceInfo:
     """Build and return the device info for EC."""
     return DeviceInfo(
         entry_type=DeviceEntryType.SERVICE,

homeassistant/components/environment_canada/camera.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import voluptuous as vol
 
 from homeassistant.components.camera import Camera
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import (
     AddEntitiesCallback,
@@ -15,7 +14,8 @@ from homeassistant.helpers.typing import VolDictType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import device_info
-from .const import ATTR_OBSERVATION_TIME, DOMAIN
+from .const import ATTR_OBSERVATION_TIME
+from .coordinator import ECConfigEntry
 
 SERVICE_SET_RADAR_TYPE = "set_radar_type"
 SET_RADAR_TYPE_SCHEMA: VolDictType = {
@@ -25,11 +25,11 @@ SET_RADAR_TYPE_SCHEMA: VolDictType = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ECConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add a weather entity from a config_entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]["radar_coordinator"]
+    coordinator = config_entry.runtime_data.radar_coordinator
     async_add_entities([ECCamera(coordinator)])
 
     platform = async_get_current_platform()

homeassistant/components/environment_canada/coordinator.py

@@ -1,29 +1,59 @@
 """Coordinator for the Environment Canada (EC) component."""
 
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import timedelta
 import logging
 import xml.etree.ElementTree as ET
 
-from env_canada import ec_exc
+from env_canada import ECAirQuality, ECRadar, ECWeather, ec_exc
 
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from .const import DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
+type ECConfigEntry = ConfigEntry[ECRuntimeData]
 
-class ECDataUpdateCoordinator(DataUpdateCoordinator):
+
+@dataclass
+class ECRuntimeData:
+    """Class to hold EC runtime data."""
+
+    aqhi_coordinator: ECDataUpdateCoordinator[ECAirQuality]
+    radar_coordinator: ECDataUpdateCoordinator[ECRadar]
+    weather_coordinator: ECDataUpdateCoordinator[ECWeather]
+
+
+class ECDataUpdateCoordinator[_ECDataTypeT: ECAirQuality | ECRadar | ECWeather](
+    DataUpdateCoordinator[_ECDataTypeT]
+):
     """Class to manage fetching EC data."""
 
-    def __init__(self, hass, ec_data, name, update_interval):
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        entry: ECConfigEntry,
+        ec_data: _ECDataTypeT,
+        name: str,
+        update_interval: timedelta,
+    ) -> None:
         """Initialize global EC data updater."""
         super().__init__(
-            hass, _LOGGER, name=f"{DOMAIN} {name}", update_interval=update_interval
+            hass,
+            _LOGGER,
+            config_entry=entry,
+            name=f"{DOMAIN} {name}",
+            update_interval=update_interval,
         )
         self.ec_data = ec_data
         self.last_update_success = False
 
-    async def _async_update_data(self):
+    async def _async_update_data(self) -> _ECDataTypeT:
         """Fetch data from EC."""
         try:
             await self.ec_data.update()

homeassistant/components/environment_canada/diagnostics.py

@@ -5,23 +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.const import CONF_LATITUDE, CONF_LONGITUDE
 from homeassistant.core import HomeAssistant
 
-from .const import DOMAIN
+from .coordinator import ECConfigEntry
 
 TO_REDACT = {CONF_LATITUDE, CONF_LONGITUDE}
 
 
 async def async_get_config_entry_diagnostics(
-    hass: HomeAssistant, config_entry: ConfigEntry
+    hass: HomeAssistant, config_entry: ECConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    coordinators = hass.data[DOMAIN][config_entry.entry_id]
-    weather_coord = coordinators["weather_coordinator"]
-
     return {
         "config_entry_data": async_redact_data(dict(config_entry.data), TO_REDACT),
-        "weather_data": dict(weather_coord.ec_data.conditions),
+        "weather_data": dict(
+            config_entry.runtime_data.weather_coordinator.ec_data.conditions
+        ),
     }

homeassistant/components/environment_canada/sensor.py

@@ -12,7 +12,6 @@ from homeassistant.components.sensor import (
     SensorEntityDescription,
     SensorStateClass,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     ATTR_LOCATION,
     DEGREE,
@@ -28,7 +27,8 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from . import device_info
-from .const import ATTR_STATION, DOMAIN
+from .const import ATTR_STATION
+from .coordinator import ECConfigEntry
 
 ATTR_TIME = "alert time"
 
@@ -251,15 +251,17 @@ ALERT_TYPES: tuple[ECSensorEntityDescription, ...] = (
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ECConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add a weather entity from a config_entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]["weather_coordinator"]
+    weather_coordinator = config_entry.runtime_data.weather_coordinator
-    sensors: list[ECBaseSensor] = [ECSensor(coordinator, desc) for desc in SENSOR_TYPES]
+    sensors: list[ECBaseSensor] = [
-    sensors.extend([ECAlertSensor(coordinator, desc) for desc in ALERT_TYPES])
+        ECSensor(weather_coordinator, desc) for desc in SENSOR_TYPES
-    aqhi_coordinator = hass.data[DOMAIN][config_entry.entry_id]["aqhi_coordinator"]
+    ]
-    sensors.append(ECSensor(aqhi_coordinator, AQHI_SENSOR))
+    sensors.extend([ECAlertSensor(weather_coordinator, desc) for desc in ALERT_TYPES])
+
+    sensors.append(ECSensor(config_entry.runtime_data.aqhi_coordinator, AQHI_SENSOR))
     async_add_entities(sensors)
 
 

homeassistant/components/environment_canada/weather.py

@@ -27,7 +27,6 @@ from homeassistant.components.weather import (
     SingleCoordinatorWeatherEntity,
     WeatherEntityFeature,
 )
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import (
     UnitOfLength,
     UnitOfPressure,
@@ -40,6 +39,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import device_info
 from .const import DOMAIN
+from .coordinator import ECConfigEntry
 
 # Icon codes from http://dd.weatheroffice.ec.gc.ca/citypage_weather/
 # docs/current_conditions_icon_code_descriptions_e.csv
@@ -61,11 +61,10 @@ ICON_CONDITION_MAP = {
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: ECConfigEntry,
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Add a weather entity from a config_entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id]["weather_coordinator"]
     entity_registry = er.async_get(hass)
 
     # Remove hourly entity from legacy config entries
@@ -76,7 +75,7 @@ async def async_setup_entry(
     ):
         entity_registry.async_remove(hourly_entity_id)
 
-    async_add_entities([ECWeather(coordinator)])
+    async_add_entities([ECWeather(config_entry.runtime_data.weather_coordinator)])
 
 
 def _calculate_unique_id(config_entry_unique_id: str | None, hourly: bool) -> str: