diff --git a/homeassistant/components/israel_rail/__init__.py b/homeassistant/components/israel_rail/__init__.py index 3c33a159a63..ed800f559d4 100644 --- a/homeassistant/components/israel_rail/__init__.py +++ b/homeassistant/components/israel_rail/__init__.py @@ -4,13 +4,12 @@ import logging from israelrailapi import TrainSchedule -from homeassistant.config_entries import ConfigEntry from homeassistant.const import Platform from homeassistant.core import HomeAssistant from homeassistant.exceptions import ConfigEntryNotReady from .const import CONF_DESTINATION, CONF_START, DOMAIN -from .coordinator import IsraelRailDataUpdateCoordinator +from .coordinator import IsraelRailConfigEntry, IsraelRailDataUpdateCoordinator _LOGGER = logging.getLogger(__name__) @@ -18,9 +17,6 @@ _LOGGER = logging.getLogger(__name__) PLATFORMS: list[Platform] = [Platform.SENSOR] -type IsraelRailConfigEntry = ConfigEntry[IsraelRailDataUpdateCoordinator] - - async def async_setup_entry(hass: HomeAssistant, entry: IsraelRailConfigEntry) -> bool: """Set up Israel rail from a config entry.""" config = entry.data @@ -43,7 +39,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: IsraelRailConfigEntry) - ) from e israel_rail_coordinator = IsraelRailDataUpdateCoordinator( - hass, train_schedule, start, destination + hass, entry, train_schedule, start, destination ) await israel_rail_coordinator.async_config_entry_first_refresh() entry.runtime_data = israel_rail_coordinator diff --git a/homeassistant/components/israel_rail/coordinator.py b/homeassistant/components/israel_rail/coordinator.py index b022e3fd790..190ed938790 100644 --- a/homeassistant/components/israel_rail/coordinator.py +++ b/homeassistant/components/israel_rail/coordinator.py @@ -38,14 +38,18 @@ def departure_time(train_route: TrainRoute) -> datetime | None: return start_datetime.astimezone() if start_datetime else None +type IsraelRailConfigEntry = ConfigEntry[IsraelRailDataUpdateCoordinator] + + class IsraelRailDataUpdateCoordinator(DataUpdateCoordinator[list[DataConnection]]): """A IsraelRail Data Update Coordinator.""" - config_entry: ConfigEntry + config_entry: IsraelRailConfigEntry def __init__( self, hass: HomeAssistant, + config_entry: IsraelRailConfigEntry, train_schedule: TrainSchedule, start: str, destination: str, @@ -54,6 +58,7 @@ class IsraelRailDataUpdateCoordinator(DataUpdateCoordinator[list[DataConnection] super().__init__( hass, _LOGGER, + config_entry=config_entry, name=DOMAIN, update_interval=DEFAULT_SCAN_INTERVAL, ) diff --git a/homeassistant/components/israel_rail/sensor.py b/homeassistant/components/israel_rail/sensor.py index 132a9a74826..d0c93da3451 100644 --- a/homeassistant/components/israel_rail/sensor.py +++ b/homeassistant/components/israel_rail/sensor.py @@ -19,9 +19,12 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback from homeassistant.helpers.typing import StateType from homeassistant.helpers.update_coordinator import CoordinatorEntity -from . import IsraelRailConfigEntry from .const import ATTRIBUTION, DEPARTURES_COUNT, DOMAIN -from .coordinator import DataConnection, IsraelRailDataUpdateCoordinator +from .coordinator import ( + DataConnection, + IsraelRailConfigEntry, + IsraelRailDataUpdateCoordinator, +) _LOGGER = logging.getLogger(__name__)