From b089f89f14f381bdf7b45880071e68b4a35f782b Mon Sep 17 00:00:00 2001 From: G Johansson Date: Wed, 1 May 2024 22:05:47 +0200 Subject: [PATCH] Use config entry runtime data in Trafikverket Ferry (#116557) --- .../components/trafikverket_ferry/__init__.py | 12 +++++------ .../trafikverket_ferry/coordinator.py | 20 +++++++++++-------- .../components/trafikverket_ferry/sensor.py | 8 +++++--- 3 files changed, 23 insertions(+), 17 deletions(-) diff --git a/homeassistant/components/trafikverket_ferry/__init__.py b/homeassistant/components/trafikverket_ferry/__init__.py index 8c8c121881f..dbcbc1a4aba 100644 --- a/homeassistant/components/trafikverket_ferry/__init__.py +++ b/homeassistant/components/trafikverket_ferry/__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 +TVFerryConfigEntry = ConfigEntry[TVDataUpdateCoordinator] -async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: + +async def async_setup_entry(hass: HomeAssistant, entry: TVFerryConfigEntry) -> bool: """Set up Trafikverket Ferry 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 Ferry config entry.""" - return await hass.config_entries.async_unload_platforms(entry, PLATFORMS) diff --git a/homeassistant/components/trafikverket_ferry/coordinator.py b/homeassistant/components/trafikverket_ferry/coordinator.py index 8d0492b1e43..cb11889345a 100644 --- a/homeassistant/components/trafikverket_ferry/coordinator.py +++ b/homeassistant/components/trafikverket_ferry/coordinator.py @@ -4,13 +4,12 @@ from __future__ import annotations from datetime import date, datetime, time, timedelta import logging -from typing import Any +from typing import TYPE_CHECKING, Any from pytrafikverket import TrafikverketFerry from pytrafikverket.exceptions import InvalidAuthentication, NoFerryFound from pytrafikverket.trafikverket_ferry import FerryStop -from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_API_KEY, CONF_WEEKDAY, WEEKDAYS from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryAuthFailed @@ -20,6 +19,9 @@ from homeassistant.util import dt as dt_util from .const import CONF_FROM, CONF_TIME, CONF_TO, DOMAIN +if TYPE_CHECKING: + from . import TVFerryConfigEntry + _LOGGER = logging.getLogger(__name__) TIME_BETWEEN_UPDATES = timedelta(minutes=5) @@ -48,7 +50,9 @@ def next_departuredate(departure: list[str]) -> date: class TVDataUpdateCoordinator(DataUpdateCoordinator): """A Trafikverket Data Update Coordinator.""" - def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None: + config_entry: TVFerryConfigEntry + + def __init__(self, hass: HomeAssistant) -> None: """Initialize the Trafikverket coordinator.""" super().__init__( hass, @@ -57,12 +61,12 @@ class TVDataUpdateCoordinator(DataUpdateCoordinator): update_interval=TIME_BETWEEN_UPDATES, ) self._ferry_api = TrafikverketFerry( - async_get_clientsession(hass), entry.data[CONF_API_KEY] + async_get_clientsession(hass), self.config_entry.data[CONF_API_KEY] ) - self._from: str = entry.data[CONF_FROM] - self._to: str = entry.data[CONF_TO] - self._time: time | None = dt_util.parse_time(entry.data[CONF_TIME]) - self._weekdays: list[str] = entry.data[CONF_WEEKDAY] + self._from: str = self.config_entry.data[CONF_FROM] + self._to: str = self.config_entry.data[CONF_TO] + self._time: time | None = dt_util.parse_time(self.config_entry.data[CONF_TIME]) + self._weekdays: list[str] = self.config_entry.data[CONF_WEEKDAY] async def _async_update_data(self) -> dict[str, Any]: """Fetch data from Trafikverket.""" diff --git a/homeassistant/components/trafikverket_ferry/sensor.py b/homeassistant/components/trafikverket_ferry/sensor.py index 93f2d1987b6..5a13159ecfd 100644 --- a/homeassistant/components/trafikverket_ferry/sensor.py +++ b/homeassistant/components/trafikverket_ferry/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 from homeassistant.core import HomeAssistant, callback from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo @@ -21,6 +20,7 @@ from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity from homeassistant.util.dt import as_utc +from . import TVFerryConfigEntry from .const import ATTRIBUTION, DOMAIN from .coordinator import TVDataUpdateCoordinator @@ -88,11 +88,13 @@ SENSOR_TYPES: tuple[TrafikverketSensorEntityDescription, ...] = ( async def async_setup_entry( - hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback + hass: HomeAssistant, + entry: TVFerryConfigEntry, + 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( [