diff --git a/homeassistant/components/airvisual/__init__.py b/homeassistant/components/airvisual/__init__.py index 60fdbf12ca1..f8f045859b3 100644 --- a/homeassistant/components/airvisual/__init__.py +++ b/homeassistant/components/airvisual/__init__.py @@ -53,6 +53,8 @@ from .const import ( LOGGER, ) +type AirVisualConfigEntry = ConfigEntry[DataUpdateCoordinator] + # We use a raw string for the airvisual_pro domain (instead of importing the actual # constant) so that we can avoid listing it as a dependency: DOMAIN_AIRVISUAL_PRO = "airvisual_pro" @@ -91,10 +93,9 @@ def async_get_cloud_coordinators_by_api_key( ) -> list[DataUpdateCoordinator]: """Get all DataUpdateCoordinator objects related to a particular API key.""" return [ - coordinator - for entry_id, coordinator in hass.data[DOMAIN].items() - if (entry := hass.config_entries.async_get_entry(entry_id)) - and entry.data.get(CONF_API_KEY) == api_key + entry.runtime_data + for entry in hass.config_entries.async_entries(DOMAIN) + if entry.data.get(CONF_API_KEY) == api_key and hasattr(entry, "runtime_data") ] @@ -172,7 +173,7 @@ def _standardize_geography_config_entry( hass.config_entries.async_update_entry(entry, **entry_updates) -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_setup_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> bool: """Set up AirVisual as config entry.""" if CONF_API_KEY not in entry.data: # If this is a migrated AirVisual Pro entry, there's no actual setup to do; @@ -220,8 +221,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: entry.async_on_unload(entry.add_update_listener(async_reload_entry)) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {}) - hass.data[DOMAIN][entry.entry_id] = coordinator + entry.runtime_data = coordinator # Reassess the interval between 2 server requests async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY]) @@ -231,7 +231,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_migrate_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> bool: """Migrate an old config entry.""" version = entry.version @@ -388,21 +388,19 @@ async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: return True -async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: +async def async_unload_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> bool: """Unload an AirVisual config entry.""" unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS) - if unload_ok: - hass.data[DOMAIN].pop(entry.entry_id) - if CONF_API_KEY in entry.data: - # Re-calculate the update interval period for any remaining consumers of - # this API key: - async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY]) + if unload_ok and CONF_API_KEY in entry.data: + # Re-calculate the update interval period for any remaining consumers of + # this API key: + async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY]) return unload_ok -async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None: +async def async_reload_entry(hass: HomeAssistant, entry: AirVisualConfigEntry) -> None: """Handle an options update.""" await hass.config_entries.async_reload(entry.entry_id) diff --git a/homeassistant/components/airvisual/diagnostics.py b/homeassistant/components/airvisual/diagnostics.py index 348bb249b0f..2e7c60364f9 100644 --- a/homeassistant/components/airvisual/diagnostics.py +++ b/homeassistant/components/airvisual/diagnostics.py @@ -5,7 +5,6 @@ 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_API_KEY, CONF_COUNTRY, @@ -15,9 +14,9 @@ from homeassistant.const import ( CONF_UNIQUE_ID, ) from homeassistant.core import HomeAssistant -from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from .const import CONF_CITY, DOMAIN +from . import AirVisualConfigEntry +from .const import CONF_CITY CONF_COORDINATES = "coordinates" CONF_TITLE = "title" @@ -37,10 +36,10 @@ TO_REDACT = { async def async_get_config_entry_diagnostics( - hass: HomeAssistant, entry: ConfigEntry + hass: HomeAssistant, entry: AirVisualConfigEntry ) -> dict[str, Any]: """Return diagnostics for a config entry.""" - coordinator: DataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data return { "entry": async_redact_data(entry.as_dict(), TO_REDACT), diff --git a/homeassistant/components/airvisual/sensor.py b/homeassistant/components/airvisual/sensor.py index df0e3da1f45..c9df2f72233 100644 --- a/homeassistant/components/airvisual/sensor.py +++ b/homeassistant/components/airvisual/sensor.py @@ -26,8 +26,8 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.update_coordinator import DataUpdateCoordinator -from . import AirVisualEntity -from .const import CONF_CITY, DOMAIN +from . import AirVisualConfigEntry, AirVisualEntity +from .const import CONF_CITY ATTR_CITY = "city" ATTR_COUNTRY = "country" @@ -105,10 +105,12 @@ POLLUTANT_UNITS = { async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirVisualConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirVisual sensors based on a config entry.""" - coordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( AirVisualGeographySensor(coordinator, entry, description, locale) for locale in GEOGRAPHY_SENSOR_LOCALES