diff --git a/homeassistant/components/environment_canada/__init__.py b/homeassistant/components/environment_canada/__init__.py index 0b6eadf6d13..c87832de6ab 100644 --- a/homeassistant/components/environment_canada/__init__.py +++ b/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( - hass, weather_data, "weather", DEFAULT_WEATHER_UPDATE_INTERVAL + weather_coordinator = ECDataUpdateCoordinator( + 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( - hass, radar_data, "radar", DEFAULT_RADAR_UPDATE_INTERVAL + radar_coordinator = ECDataUpdateCoordinator( + 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( - hass, aqhi_data, "AQHI", DEFAULT_WEATHER_UPDATE_INTERVAL + aqhi_coordinator = ECDataUpdateCoordinator( + 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, {}) - hass.data[DOMAIN][config_entry.entry_id] = coordinators + config_entry.runtime_data = ECRuntimeData( + 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( - config_entry, PLATFORMS - ) - - hass.data[DOMAIN].pop(config_entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(config_entry, PLATFORMS) -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, diff --git a/homeassistant/components/environment_canada/camera.py b/homeassistant/components/environment_canada/camera.py index 1625cd253da..d0497d855e5 100644 --- a/homeassistant/components/environment_canada/camera.py +++ b/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() diff --git a/homeassistant/components/environment_canada/coordinator.py b/homeassistant/components/environment_canada/coordinator.py index 8e77b309c78..8161e26028c 100644 --- a/homeassistant/components/environment_canada/coordinator.py +++ b/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() diff --git a/homeassistant/components/environment_canada/diagnostics.py b/homeassistant/components/environment_canada/diagnostics.py index 0fb565fda59..024cca15f12 100644 --- a/homeassistant/components/environment_canada/diagnostics.py +++ b/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 + ), } diff --git a/homeassistant/components/environment_canada/sensor.py b/homeassistant/components/environment_canada/sensor.py index 1a5d096203d..ddececa8132 100644 --- a/homeassistant/components/environment_canada/sensor.py +++ b/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"] - sensors: list[ECBaseSensor] = [ECSensor(coordinator, desc) for desc in SENSOR_TYPES] - sensors.extend([ECAlertSensor(coordinator, desc) for desc in ALERT_TYPES]) - aqhi_coordinator = hass.data[DOMAIN][config_entry.entry_id]["aqhi_coordinator"] - sensors.append(ECSensor(aqhi_coordinator, AQHI_SENSOR)) + weather_coordinator = config_entry.runtime_data.weather_coordinator + sensors: list[ECBaseSensor] = [ + ECSensor(weather_coordinator, desc) for desc in SENSOR_TYPES + ] + 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) diff --git a/homeassistant/components/environment_canada/weather.py b/homeassistant/components/environment_canada/weather.py index 1871062c2e9..e49164d6b81 100644 --- a/homeassistant/components/environment_canada/weather.py +++ b/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: