diff --git a/homeassistant/components/airgradient/__init__.py b/homeassistant/components/airgradient/__init__.py index da3edcf0453..91ee0a440a6 100644 --- a/homeassistant/components/airgradient/__init__.py +++ b/homeassistant/components/airgradient/__init__.py @@ -2,6 +2,8 @@ from __future__ import annotations +from dataclasses import dataclass + from airgradient import AirGradientClient from homeassistant.config_entries import ConfigEntry @@ -16,6 +18,17 @@ from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoo PLATFORMS: list[Platform] = [Platform.SELECT, Platform.SENSOR] +@dataclass +class AirGradientData: + """AirGradient data class.""" + + measurement: AirGradientMeasurementCoordinator + config: AirGradientConfigCoordinator + + +type AirGradientConfigEntry = ConfigEntry[AirGradientData] + + async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up Airgradient from a config entry.""" @@ -39,10 +52,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: sw_version=measurement_coordinator.data.firmware_version, ) - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = { - "measurement": measurement_coordinator, - "config": config_coordinator, - } + entry.runtime_data = AirGradientData( + measurement=measurement_coordinator, + config=config_coordinator, + ) await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) @@ -51,7 +64,4 @@ 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): - hass.data[DOMAIN].pop(entry.entry_id) - - return unload_ok + return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/airgradient/coordinator.py b/homeassistant/components/airgradient/coordinator.py index 90aded9a4ba..fbc1505f9c3 100644 --- a/homeassistant/components/airgradient/coordinator.py +++ b/homeassistant/components/airgradient/coordinator.py @@ -1,21 +1,26 @@ """Define an object to manage fetching AirGradient data.""" +from __future__ import annotations + from datetime import timedelta +from typing import TYPE_CHECKING from airgradient import AirGradientClient, AirGradientError, Config, Measures -from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from .const import LOGGER +if TYPE_CHECKING: + from . import AirGradientConfigEntry + class AirGradientCoordinator[_DataT](DataUpdateCoordinator[_DataT]): """Class to manage fetching AirGradient data.""" _update_interval: timedelta - config_entry: ConfigEntry + config_entry: AirGradientConfigEntry def __init__(self, hass: HomeAssistant, client: AirGradientClient) -> None: """Initialize coordinator.""" diff --git a/homeassistant/components/airgradient/select.py b/homeassistant/components/airgradient/select.py index 7a82d3b8a46..7880e55de19 100644 --- a/homeassistant/components/airgradient/select.py +++ b/homeassistant/components/airgradient/select.py @@ -7,14 +7,14 @@ from airgradient import AirGradientClient, Config from airgradient.models import ConfigurationControl, TemperatureUnit from homeassistant.components.select import SelectEntity, SelectEntityDescription -from homeassistant.config_entries import ConfigEntry from homeassistant.const import EntityCategory from homeassistant.core import HomeAssistant from homeassistant.exceptions import ServiceValidationError from homeassistant.helpers.entity_platform import AddEntitiesCallback +from . import AirGradientConfigEntry from .const import DOMAIN -from .coordinator import AirGradientConfigCoordinator, AirGradientMeasurementCoordinator +from .coordinator import AirGradientConfigCoordinator from .entity import AirGradientEntity @@ -56,16 +56,14 @@ PROTECTED_SELECT_TYPES: tuple[AirGradientSelectEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirGradientConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirGradient select entities based on a config entry.""" - config_coordinator: AirGradientConfigCoordinator = hass.data[DOMAIN][ - entry.entry_id - ]["config"] - measurement_coordinator: AirGradientMeasurementCoordinator = hass.data[DOMAIN][ - entry.entry_id - ]["measurement"] + config_coordinator = entry.runtime_data.config + measurement_coordinator = entry.runtime_data.measurement entities = [AirGradientSelect(config_coordinator, CONFIG_CONTROL_ENTITY)] diff --git a/homeassistant/components/airgradient/sensor.py b/homeassistant/components/airgradient/sensor.py index f21f13b80ab..6123d4289f9 100644 --- a/homeassistant/components/airgradient/sensor.py +++ b/homeassistant/components/airgradient/sensor.py @@ -11,7 +11,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( CONCENTRATION_MICROGRAMS_PER_CUBIC_METER, CONCENTRATION_PARTS_PER_MILLION, @@ -24,7 +23,7 @@ from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType -from .const import DOMAIN +from . import AirGradientConfigEntry from .coordinator import AirGradientMeasurementCoordinator from .entity import AirGradientEntity @@ -127,13 +126,13 @@ SENSOR_TYPES: tuple[AirGradientSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: AirGradientConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up AirGradient sensor entities based on a config entry.""" - coordinator: AirGradientMeasurementCoordinator = hass.data[DOMAIN][entry.entry_id][ - "measurement" - ] + coordinator = entry.runtime_data.measurement listener: Callable[[], None] | None = None not_setup: set[AirGradientSensorEntityDescription] = set(SENSOR_TYPES)