From 3bf67f3dddce45f747bf3d784f6111066b978f91 Mon Sep 17 00:00:00 2001 From: G Johansson Date: Wed, 1 May 2024 22:08:19 +0200 Subject: [PATCH] Use config entry runtime data in Trafikverket Train (#116556) --- .../components/trafikverket_train/__init__.py | 12 +++++------ .../trafikverket_train/coordinator.py | 21 ++++++++++++------- .../components/trafikverket_train/sensor.py | 8 ++++--- 3 files changed, 24 insertions(+), 17 deletions(-) diff --git a/homeassistant/components/trafikverket_train/__init__.py b/homeassistant/components/trafikverket_train/__init__.py index 8b427c3431d..4bf1f681807 100644 --- a/homeassistant/components/trafikverket_train/__init__.py +++ b/homeassistant/components/trafikverket_train/__init__.py @@ -16,11 +16,13 @@ from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady from homeassistant.helpers import entity_registry as er from homeassistant.helpers.aiohttp_client import async_get_clientsession -from .const import CONF_FILTER_PRODUCT, CONF_FROM, CONF_TO, DOMAIN, PLATFORMS +from .const import CONF_FROM, CONF_TO, PLATFORMS from .coordinator import TVDataUpdateCoordinator +TVTrainConfigEntry = ConfigEntry[TVDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: TVTrainConfigEntry) -> bool: """Set up Trafikverket Train from a config entry.""" http_session = async_get_clientsession(hass) @@ -37,11 +39,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: f" {entry.data[CONF_TO]}. Error: {error} " ) from error - coordinator = TVDataUpdateCoordinator( - hass, entry, to_station, from_station, entry.options.get(CONF_FILTER_PRODUCT) - ) + coordinator = TVDataUpdateCoordinator(hass, to_station, from_station) await coordinator.async_config_entry_first_refresh() - hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator + entry.runtime_data = coordinator entity_reg = er.async_get(hass) entries = er.async_entries_for_config_entry(entity_reg, entry.entry_id) diff --git a/homeassistant/components/trafikverket_train/coordinator.py b/homeassistant/components/trafikverket_train/coordinator.py index cf78228ed58..e56f5d3a2e9 100644 --- a/homeassistant/components/trafikverket_train/coordinator.py +++ b/homeassistant/components/trafikverket_train/coordinator.py @@ -5,6 +5,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime, time, timedelta import logging +from typing import TYPE_CHECKING from pytrafikverket import TrafikverketTrain from pytrafikverket.exceptions import ( @@ -14,7 +15,6 @@ from pytrafikverket.exceptions import ( ) from pytrafikverket.trafikverket_train import StationInfo, TrainStop -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_WEEKDAY from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed @@ -22,9 +22,12 @@ from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed from homeassistant.util import dt as dt_util -from .const import CONF_TIME, DOMAIN +from .const import CONF_FILTER_PRODUCT, CONF_TIME, DOMAIN from .util import next_departuredate +if TYPE_CHECKING: + from . import TVTrainConfigEntry + @dataclass class TrainData: @@ -65,13 +68,13 @@ def _get_as_joined(information: list[str] | None) -> str | None: class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]): """A Trafikverket Data Update Coordinator.""" + config_entry: TVTrainConfigEntry + def __init__( self, hass: HomeAssistant, - entry: ConfigEntry, to_station: StationInfo, from_station: StationInfo, - filter_product: str | None, ) -> None: """Initialize the Trafikverket coordinator.""" super().__init__( @@ -81,13 +84,15 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator[TrainData]): update_interval=TIME_BETWEEN_UPDATES, ) self._train_api = TrafikverketTrain( - async_get_clientsession(hass), entry.data[CONF_API_KEY] + async_get_clientsession(hass), self.config_entry.data[CONF_API_KEY] ) self.from_station: StationInfo = from_station self.to_station: StationInfo = to_station - self._time: time | None = dt_util.parse_time(entry.data[CONF_TIME]) - self._weekdays: list[str] = entry.data[CONF_WEEKDAY] - self._filter_product = filter_product + self._time: time | None = dt_util.parse_time(self.config_entry.data[CONF_TIME]) + self._weekdays: list[str] = self.config_entry.data[CONF_WEEKDAY] + self._filter_product: str | None = self.config_entry.options.get( + CONF_FILTER_PRODUCT + ) async def _async_update_data(self) -> TrainData: """Fetch data from Trafikverket.""" diff --git a/homeassistant/components/trafikverket_train/sensor.py b/homeassistant/components/trafikverket_train/sensor.py index 22d8aba4725..e5331a47d16 100644 --- a/homeassistant/components/trafikverket_train/sensor.py +++ b/homeassistant/components/trafikverket_train/sensor.py @@ -12,7 +12,6 @@ from homeassistant.components.sensor import ( SensorEntity, SensorEntityDescription, ) -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_NAME, UnitOfTime from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo @@ -20,6 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity +from . import TVTrainConfigEntry from .const import ATTRIBUTION, DOMAIN from .coordinator import TrainData, TVDataUpdateCoordinator @@ -106,11 +106,13 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: TVTrainConfigEntry, + 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( [