Use config entry runtime data in Trafikverket Train (#116556)

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)

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."""

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(
         [