sensor.py

Go to the documentation of this file.

"""Get ride details and liveboard details for NMBS (Belgian railway)."""
 
from __future__ import annotations
 
import logging
 
from pyrail import iRail
import voluptuous as vol
 
from homeassistant.components.sensor import (
    PLATFORM_SCHEMA as SENSOR_PLATFORM_SCHEMA,
    SensorEntity,
)
from homeassistant.const import (
    ATTR_LATITUDE,
    ATTR_LONGITUDE,
    CONF_NAME,
    CONF_SHOW_ON_MAP,
    UnitOfTime,
)
from homeassistant.core import HomeAssistant
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.typing import ConfigType, DiscoveryInfoType
import homeassistant.util.dt as dt_util
 
_LOGGER = logging.getLogger(__name__)
 
API_FAILURE = -1
 
DEFAULT_NAME = "NMBS"
 
DEFAULT_ICON = "mdi:train"
DEFAULT_ICON_ALERT = "mdi:alert-octagon"
 
CONF_STATION_FROM = "station_from"
CONF_STATION_TO = "station_to"
CONF_STATION_LIVE = "station_live"
CONF_EXCLUDE_VIAS = "exclude_vias"
 
PLATFORM_SCHEMA = SENSOR_PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_STATION_FROM): cv.string,
        vol.Required(CONF_STATION_TO): cv.string,
        vol.Optional(CONF_STATION_LIVE): cv.string,
        vol.Optional(CONF_EXCLUDE_VIAS, default=False): cv.boolean,
        vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
        vol.Optional(CONF_SHOW_ON_MAP, default=False): cv.boolean,
    }
)
 
 
def get_time_until(departure_time=None):
    """Calculate the time between now and a train's departure time."""
    if departure_time is None:
        return 0
 
    delta = dt_util.utc_from_timestamp(int(departure_time)) - dt_util.now()
    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, arrival_time, delay=0):
    """Calculate the total travel time in minutes."""
    duration = dt_util.utc_from_timestamp(
        int(arrival_time)
    ) - dt_util.utc_from_timestamp(int(departure_time))
    duration_time = int(round(duration.total_seconds() / 60))
    return duration_time + get_delay_in_minutes(delay)
 
 
def setup_platform(
    hass: HomeAssistant,
    config: ConfigType,
    add_entities: AddEntitiesCallback,
    discovery_info: DiscoveryInfoType | None = None,
) -> None:
    """Set up the NMBS sensor with iRail API."""
 
    api_client = iRail()
 
    name = config[CONF_NAME]
    show_on_map = config[CONF_SHOW_ON_MAP]
    station_from = config[CONF_STATION_FROM]
    station_to = config[CONF_STATION_TO]
    station_live = config.get(CONF_STATION_LIVE)
    excl_vias = config[CONF_EXCLUDE_VIAS]
 
    sensors: list[SensorEntity] = [
        NMBSSensor(api_client, name, show_on_map, station_from, station_to, excl_vias)
    ]
 
    if station_live is not None:
        sensors.append(
            NMBSLiveBoard(api_client, station_live, station_from, station_to)
        )
 
    add_entities(sensors, True)
 
 
class NMBSLiveBoard(SensorEntity):
    """Get the next train from a station's liveboard."""
 
    _attr_attribution = "https://api.irail.be/"
 
    def __init__(self, api_client, live_station, station_from, station_to):
        """Initialize the sensor for getting liveboard data."""
        self._station_station = live_station
        self._api_client_api_client = api_client
        self._station_from_station_from = station_from
        self._station_to_station_to = station_to
        self._attrs_attrs = {}
        self._state_state = None
 
    @property
    def name(self):
        """Return the sensor default name."""
        return f"NMBS Live ({self._station})"
 
    @property
    def unique_id(self):
        """Return a unique ID."""
        unique_id = f"{self._station}_{self._station_from}_{self._station_to}"
 
        return f"nmbs_live_{unique_id}"
 
    @property
    def icon(self):
        """Return the default icon or an alert icon if delays."""
        if self._attrs_attrs and int(self._attrs_attrs["delay"]) > 0:
            return DEFAULT_ICON_ALERT
 
        return DEFAULT_ICON
 
    @property
    def native_value(self):
        """Return sensor state."""
        return self._state_state
 
    @property
    def extra_state_attributes(self):
        """Return the sensor attributes if data is available."""
        if self._state_state is None or not self._attrs_attrs:
            return None
 
        delay = get_delay_in_minutes(self._attrs_attrs["delay"])
        departure = get_time_until(self._attrs_attrs["time"])
 
        attrs = {
            "departure": f"In {departure} minutes",
            "departure_minutes": departure,
            "extra_train": int(self._attrs_attrs["isExtra"]) > 0,
            "vehicle_id": self._attrs_attrs["vehicle"],
            "monitored_station": self._station_station,
        }
 
        if delay > 0:
            attrs["delay"] = f"{delay} minutes"
            attrs["delay_minutes"] = delay
 
        return attrs
 
    def update(self) -> None:
        """Set the state equal to the next departure."""
        liveboard = self._api_client_api_client.get_liveboard(self._station_station)
 
        if liveboard == API_FAILURE:
            _LOGGER.warning("API failed in NMBSLiveBoard")
            return
 
        if not (departures := liveboard.get("departures")):
            _LOGGER.warning("API returned invalid departures: %r", liveboard)
            return
 
        _LOGGER.debug("API returned departures: %r", departures)
        if departures["number"] == "0":
            # No trains are scheduled
            return
        next_departure = departures["departure"][0]
 
        self._attrs_attrs = next_departure
        self._state_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, name, show_on_map, station_from, station_to, excl_vias
    ):
        """Initialize the NMBS connection sensor."""
        self._name_name = name
        self._show_on_map_show_on_map = show_on_map
        self._api_client_api_client = api_client
        self._station_from_station_from = station_from
        self._station_to_station_to = station_to
        self._excl_vias_excl_vias = excl_vias
 
        self._attrs_attrs = {}
        self._state_state = None
 
    @property
    def name(self):
        """Return the name of the sensor."""
        return self._name_name
 
    @property
    def icon(self):
        """Return the sensor default icon or an alert icon if any delay."""
        if self._attrs_attrs:
            delay = get_delay_in_minutes(self._attrs_attrs["departure"]["delay"])
            if delay > 0:
                return "mdi:alert-octagon"
 
        return "mdi:train"
 
    @property
    def extra_state_attributes(self):
        """Return sensor attributes if data is available."""
        if self._state_state is None or not self._attrs_attrs:
            return None
 
        delay = get_delay_in_minutes(self._attrs_attrs["departure"]["delay"])
        departure = get_time_until(self._attrs_attrs["departure"]["time"])
        canceled = int(self._attrs_attrs["departure"]["canceled"])
 
        attrs = {
            "destination": self._station_to_station_to,
            "direction": self._attrs_attrs["departure"]["direction"]["name"],
            "platform_arriving": self._attrs_attrs["arrival"]["platform"],
            "platform_departing": self._attrs_attrs["departure"]["platform"],
            "vehicle_id": self._attrs_attrs["departure"]["vehicle"],
        }
 
        if canceled != 1:
            attrs["departure"] = f"In {departure} minutes"
            attrs["departure_minutes"] = departure
            attrs["canceled"] = False
        else:
            attrs["departure"] = None
            attrs["departure_minutes"] = None
            attrs["canceled"] = True
 
        if self._show_on_map_show_on_map and self.station_coordinatesstation_coordinates:
            attrs[ATTR_LATITUDE] = self.station_coordinatesstation_coordinates[0]
            attrs[ATTR_LONGITUDE] = self.station_coordinatesstation_coordinates[1]
 
        if self.is_via_connectionis_via_connection and not self._excl_vias_excl_vias:
            via = self._attrs_attrs["vias"]["via"][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"])
 
        if delay > 0:
            attrs["delay"] = f"{delay} minutes"
            attrs["delay_minutes"] = delay
 
        return attrs
 
    @property
    def native_value(self):
        """Return the state of the device."""
        return self._state_state
 
    @property
    def station_coordinates(self):
        """Get the lat, long coordinates for station."""
        if self._state_state is None or not self._attrs_attrs:
            return []
 
        latitude = float(self._attrs_attrs["departure"]["stationinfo"]["locationY"])
        longitude = float(self._attrs_attrs["departure"]["stationinfo"]["locationX"])
        return [latitude, longitude]
 
    @property
    def is_via_connection(self):
        """Return whether the connection goes through another station."""
        if not self._attrs_attrs:
            return False
 
        return "vias" in self._attrs_attrs and int(self._attrs_attrs["vias"]["number"]) > 0
 
    def update(self) -> None:
        """Set the state to the duration of a connection."""
        connections = self._api_client_api_client.get_connections(
            self._station_from_station_from, self._station_to_station_to
        )
 
        if connections == API_FAILURE:
            _LOGGER.warning("API failed in NMBSSensor")
            return
 
        if not (connection := connections.get("connection")):
            _LOGGER.warning("API returned invalid connection: %r", connections)
            return
 
        _LOGGER.debug("API returned connection: %r", connection)
        if int(connection[0]["departure"]["left"]) > 0:
            next_connection = connection[1]
        else:
            next_connection = connection[0]
 
        self._attrs_attrs = next_connection
 
        if self._excl_vias_excl_vias and self.is_via_connectionis_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_state = duration