diff --git a/homeassistant/components/trafikverket_weatherstation/__init__.py b/homeassistant/components/trafikverket_weatherstation/__init__.py index e1cd9c90909..1bd7fc69ae4 100644 --- a/homeassistant/components/trafikverket_weatherstation/__init__.py +++ b/homeassistant/components/trafikverket_weatherstation/__init__.py @@ -5,17 +5,18 @@ from __future__ import annotations from homeassistant.config_entries import ConfigEntry from homeassistant.core import HomeAssistant -from .const import DOMAIN, PLATFORMS +from .const import PLATFORMS from .coordinator import TVDataUpdateCoordinator +TVWeatherConfigEntry = ConfigEntry[TVDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: TVWeatherConfigEntry) -> bool: """Set up Trafikverket Weatherstation from a config entry.""" - coordinator = TVDataUpdateCoordinator(hass, entry) + coordinator = TVDataUpdateCoordinator(hass) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator - + entry.runtime_data = coordinator await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS) return True @@ -23,5 +24,4 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Unload Trafikverket Weatherstation config entry.""" - return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/trafikverket_weatherstation/coordinator.py b/homeassistant/components/trafikverket_weatherstation/coordinator.py index 508ae7eec16..e0319b1b932 100644 --- a/homeassistant/components/trafikverket_weatherstation/coordinator.py +++ b/homeassistant/components/trafikverket_weatherstation/coordinator.py @@ -4,6 +4,7 @@ from __future__ import annotations from datetime import timedelta import logging +from typing import TYPE_CHECKING from pytrafikverket.exceptions import ( InvalidAuthentication, @@ -12,7 +13,6 @@ from pytrafikverket.exceptions import ( ) from pytrafikverket.trafikverket_weather import TrafikverketWeather, WeatherStationInfo -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed @@ -21,6 +21,9 @@ from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, Upda from .const import CONF_STATION, DOMAIN +if TYPE_CHECKING: + from . import TVWeatherConfigEntry + _LOGGER = logging.getLogger(__name__) TIME_BETWEEN_UPDATES = timedelta(minutes=10) @@ -28,7 +31,9 @@ TIME_BETWEEN_UPDATES = timedelta(minutes=10) class TVDataUpdateCoordinator(DataUpdateCoordinator[WeatherStationInfo]): """A Sensibo Data Update Coordinator.""" - def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None: + config_entry: TVWeatherConfigEntry + + def __init__(self, hass: HomeAssistant) -> None: """Initialize the Sensibo coordinator.""" super().__init__( hass, @@ -37,9 +42,9 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[WeatherStationInfo]): update_interval=TIME_BETWEEN_UPDATES, ) self._weather_api = TrafikverketWeather( - async_get_clientsession(hass), entry.data[CONF_API_KEY] + async_get_clientsession(hass), self.config_entry.data[CONF_API_KEY] ) - self._station = entry.data[CONF_STATION] + self._station = self.config_entry.data[CONF_STATION] async def _async_update_data(self) -> WeatherStationInfo: """Fetch data from Trafikverket.""" diff --git a/homeassistant/components/trafikverket_weatherstation/sensor.py b/homeassistant/components/trafikverket_weatherstation/sensor.py index bd15c34ff01..4bd14448546 100644 --- a/homeassistant/components/trafikverket_weatherstation/sensor.py +++ b/homeassistant/components/trafikverket_weatherstation/sensor.py @@ -14,7 +14,6 @@ from homeassistant.components.sensor import ( SensorEntityDescription, SensorStateClass, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( DEGREE, PERCENTAGE, @@ -30,6 +29,7 @@ from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util import dt as dt_util +from . import TVWeatherConfigEntry from .const import ATTRIBUTION, CONF_STATION, DOMAIN from .coordinator import TVDataUpdateCoordinator @@ -200,11 +200,13 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: TVWeatherConfigEntry, + async_add_entities: AddEntitiesCallback, ) -> None: """Set up the Trafikverket sensor entry.""" - coordinator: TVDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id] + coordinator = entry.runtime_data async_add_entities( TrafikverketWeatherStation(