diff --git a/homeassistant/components/vasttrafik/sensor.py b/homeassistant/components/vasttrafik/sensor.py index 2b904faab4d..8b1609be6ba 100644 --- a/homeassistant/components/vasttrafik/sensor.py +++ b/homeassistant/components/vasttrafik/sensor.py @@ -80,14 +80,23 @@ class VasttrafikDepartureSensor(Entity): """Initialize the sensor.""" self._planner = planner self._name = name or departure - self._departure = planner.location_name(departure)[0] - self._heading = planner.location_name(heading)[0] if heading else None + self._departure = self.get_station_id(departure) + self._heading = self.get_station_id(heading) if heading else None self._lines = lines if lines else None self._delay = timedelta(minutes=delay) self._departureboard = None self._state = None self._attributes = None + def get_station_id(self, location): + """Get the station ID.""" + if location.isdecimal(): + station_info = {"station_name": location, "station_id": location} + else: + station_id = self._planner.location_name(location)[0]["id"] + station_info = {"station_name": location, "station_id": station_id} + return station_info + @property def name(self): """Return the name of the sensor.""" @@ -113,8 +122,8 @@ class VasttrafikDepartureSensor(Entity): """Get the departure board.""" try: self._departureboard = self._planner.departureboard( - self._departure["id"], - direction=self._heading["id"] if self._heading else None, + self._departure["station_id"], + direction=self._heading["station_id"] if self._heading else None, date=now() + self._delay, ) except vasttrafik.Error: @@ -123,9 +132,9 @@ class VasttrafikDepartureSensor(Entity): if not self._departureboard: _LOGGER.debug( - "No departures from %s heading %s", - self._departure["name"], - self._heading["name"] if self._heading else "ANY", + "No departures from departure station %s " "to destination station %s", + self._departure["station_name"], + self._heading["station_name"] if self._heading else "ANY", ) self._state = None self._attributes = {}