"""Get ride details and liveboard details for NMBS (Belgian railway).""" from __future__ import annotations from datetime import datetime import logging from typing import Any from pyrail import iRail from pyrail.models import ConnectionDetails, LiveboardDeparture, StationDetails from homeassistant.components.sensor import SensorEntity from homeassistant.config_entries import ConfigEntry from homeassistant.const import ( ATTR_LATITUDE, ATTR_LONGITUDE, CONF_NAME, CONF_SHOW_ON_MAP, UnitOfTime, ) from homeassistant.core import HomeAssistant from homeassistant.helpers.aiohttp_client import async_get_clientsession from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback from homeassistant.util import dt as dt_util from .const import ( # noqa: F401 CONF_EXCLUDE_VIAS, CONF_STATION_FROM, CONF_STATION_LIVE, CONF_STATION_TO, DOMAIN, PLATFORMS, find_station, find_station_by_name, ) _LOGGER = logging.getLogger(__name__) DEFAULT_ICON = "mdi:train" DEFAULT_ICON_ALERT = "mdi:alert-octagon" def get_time_until(departure_time: datetime | None = None): """Calculate the time between now and a train's departure time.""" if departure_time is None: return 0 delta = dt_util.as_utc(departure_time) - dt_util.utcnow() return round(delta.total_seconds() / 60) def get_delay_in_minutes(delay=0): """Get the delay in minutes from a delay in seconds.""" return round(int(delay) / 60) def get_ride_duration(departure_time: datetime, arrival_time: datetime, delay=0): """Calculate the total travel time in minutes.""" duration = arrival_time - departure_time duration_time = int(round(duration.total_seconds() / 60)) return duration_time + get_delay_in_minutes(delay) async def async_setup_entry( hass: HomeAssistant, config_entry: ConfigEntry, async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up NMBS sensor entities based on a config entry.""" api_client = iRail(session=async_get_clientsession(hass)) name = config_entry.data.get(CONF_NAME, None) show_on_map = config_entry.data.get(CONF_SHOW_ON_MAP, False) excl_vias = config_entry.data.get(CONF_EXCLUDE_VIAS, False) station_from = find_station(hass, config_entry.data[CONF_STATION_FROM]) station_to = find_station(hass, config_entry.data[CONF_STATION_TO]) # setup the connection from station to station # setup a disabled liveboard for both from and to station async_add_entities( [ NMBSSensor( api_client, name, show_on_map, station_from, station_to, excl_vias ), NMBSLiveBoard( api_client, station_from, station_from, station_to, excl_vias ), NMBSLiveBoard(api_client, station_to, station_from, station_to, excl_vias), ] ) class NMBSLiveBoard(SensorEntity): """Get the next train from a station's liveboard.""" _attr_attribution = "https://api.irail.be/" def __init__( self, api_client: iRail, live_station: StationDetails, station_from: StationDetails, station_to: StationDetails, excl_vias: bool, ) -> None: """Initialize the sensor for getting liveboard data.""" self._station = live_station self._api_client = api_client self._station_from = station_from self._station_to = station_to self._excl_vias = excl_vias self._attrs: LiveboardDeparture | None = None self._state: str | None = None self.entity_registry_enabled_default = False @property def name(self) -> str: """Return the sensor default name.""" return f"Trains in {self._station.standard_name}" @property def unique_id(self) -> str: """Return the unique ID.""" unique_id = f"{self._station.id}_{self._station_from.id}_{self._station_to.id}" vias = "_excl_vias" if self._excl_vias else "" return f"nmbs_live_{unique_id}{vias}" @property def icon(self) -> str: """Return the default icon or an alert icon if delays.""" if self._attrs and int(self._attrs.delay) > 0: return DEFAULT_ICON_ALERT return DEFAULT_ICON @property def native_value(self) -> str | None: """Return sensor state.""" return self._state @property def extra_state_attributes(self) -> dict[str, Any] | None: """Return the sensor attributes if data is available.""" if self._state is None or not self._attrs: return None delay = get_delay_in_minutes(self._attrs.delay) departure = get_time_until(self._attrs.time) attrs = { "departure": f"In {departure} minutes", "departure_minutes": departure, "extra_train": self._attrs.is_extra, "vehicle_id": self._attrs.vehicle, "monitored_station": self._station.standard_name, } if delay > 0: attrs["delay"] = f"{delay} minutes" attrs["delay_minutes"] = delay return attrs async def async_update(self, **kwargs: Any) -> None: """Set the state equal to the next departure.""" liveboard = await self._api_client.get_liveboard(self._station.id) if liveboard is None: _LOGGER.warning("API failed in NMBSLiveBoard") return if not (departures := liveboard.departures): _LOGGER.warning("API returned invalid departures: %r", liveboard) return _LOGGER.debug("API returned departures: %r", departures) if len(departures) == 0: # No trains are scheduled return next_departure = departures[0] self._attrs = next_departure self._state = f"Track {next_departure.platform} - {next_departure.station}" class NMBSSensor(SensorEntity): """Get the total travel time for a given connection.""" _attr_attribution = "https://api.irail.be/" _attr_native_unit_of_measurement = UnitOfTime.MINUTES def __init__( self, api_client: iRail, name: str, show_on_map: bool, station_from: StationDetails, station_to: StationDetails, excl_vias: bool, ) -> None: """Initialize the NMBS connection sensor.""" self._name = name self._show_on_map = show_on_map self._api_client = api_client self._station_from = station_from self._station_to = station_to self._excl_vias = excl_vias self._attrs: ConnectionDetails | None = None self._state = None @property def unique_id(self) -> str: """Return the unique ID.""" unique_id = f"{self._station_from.id}_{self._station_to.id}" vias = "_excl_vias" if self._excl_vias else "" return f"nmbs_connection_{unique_id}{vias}" @property def name(self) -> str: """Return the name of the sensor.""" if self._name is None: return f"Train from {self._station_from.standard_name} to {self._station_to.standard_name}" return self._name @property def icon(self) -> str: """Return the sensor default icon or an alert icon if any delay.""" if self._attrs: delay = get_delay_in_minutes(self._attrs.departure.delay) if delay > 0: return "mdi:alert-octagon" return "mdi:train" @property def extra_state_attributes(self) -> dict[str, Any] | None: """Return sensor attributes if data is available.""" if self._state is None or not self._attrs: return None delay = get_delay_in_minutes(self._attrs.departure.delay) departure = get_time_until(self._attrs.departure.time) attrs = { "destination": self._attrs.departure.station, "direction": self._attrs.departure.direction.name, "platform_arriving": self._attrs.arrival.platform, "platform_departing": self._attrs.departure.platform, "vehicle_id": self._attrs.departure.vehicle, } attrs["canceled"] = self._attrs.departure.canceled if attrs["canceled"]: attrs["departure"] = None attrs["departure_minutes"] = None else: attrs["departure"] = f"In {departure} minutes" attrs["departure_minutes"] = departure if self._show_on_map and self.station_coordinates: attrs[ATTR_LATITUDE] = self.station_coordinates[0] attrs[ATTR_LONGITUDE] = self.station_coordinates[1] if self.is_via_connection and not self._excl_vias: via = self._attrs.vias[0] attrs["via"] = via.station attrs["via_arrival_platform"] = via.arrival.platform attrs["via_transfer_platform"] = via.departure.platform attrs["via_transfer_time"] = get_delay_in_minutes( via.timebetween ) + get_delay_in_minutes(via.departure.delay) attrs["delay"] = f"{delay} minutes" attrs["delay_minutes"] = delay return attrs @property def native_value(self) -> int | None: """Return the state of the device.""" return self._state @property def station_coordinates(self) -> list[float]: """Get the lat, long coordinates for station.""" if self._state is None or not self._attrs: return [] latitude = float(self._attrs.departure.station_info.latitude) longitude = float(self._attrs.departure.station_info.longitude) return [latitude, longitude] @property def is_via_connection(self) -> bool: """Return whether the connection goes through another station.""" if not self._attrs: return False return self._attrs.vias is not None and len(self._attrs.vias) > 0 async def async_update(self, **kwargs: Any) -> None: """Set the state to the duration of a connection.""" connections = await self._api_client.get_connections( self._station_from.id, self._station_to.id ) if connections is None: _LOGGER.warning("API failed in NMBSSensor") return if not (connection := connections.connections): _LOGGER.warning("API returned invalid connection: %r", connections) return _LOGGER.debug("API returned connection: %r", connection) if connection[0].departure.left: next_connection = connection[1] else: next_connection = connection[0] self._attrs = next_connection if self._excl_vias and self.is_via_connection: _LOGGER.debug( "Skipping update of NMBSSensor because this connection is a via" ) return duration = get_ride_duration( next_connection.departure.time, next_connection.arrival.time, next_connection.departure.delay, ) self._state = duration