diff --git a/homeassistant/components/swiss_public_transport/coordinator.py b/homeassistant/components/swiss_public_transport/coordinator.py index 93b3312b09916d..97253d5776e586 100644 --- a/homeassistant/components/swiss_public_transport/coordinator.py +++ b/homeassistant/components/swiss_public_transport/coordinator.py @@ -1,7 +1,7 @@ """DataUpdateCoordinator for the swiss_public_transport integration.""" from __future__ import annotations -from datetime import timedelta +from datetime import datetime, timedelta import logging from typing import TypedDict @@ -21,9 +21,9 @@ class DataConnection(TypedDict): """A connection data class.""" - departure: str - next_departure: str - next_on_departure: str + departure: datetime | None + next_departure: str | None + next_on_departure: str | None duration: str platform: str remaining_time: str @@ -58,18 +58,35 @@ async def _async_update_data(self) -> DataConnection: ) raise UpdateFailed from e - departure_time = dt_util.parse_datetime( - self._opendata.connections[0]["departure"] + departure_time = ( + dt_util.parse_datetime(self._opendata.connections[0]["departure"]) + if self._opendata.connections[0] is not None + else None ) + next_departure_time = ( + dt_util.parse_datetime(self._opendata.connections[1]["departure"]) + if self._opendata.connections[1] is not None + else None + ) + next_on_departure_time = ( + dt_util.parse_datetime(self._opendata.connections[2]["departure"]) + if self._opendata.connections[2] is not None + else None + ) + if departure_time: remaining_time = departure_time - dt_util.as_local(dt_util.utcnow()) else: remaining_time = None return DataConnection( - departure=self._opendata.connections[0]["departure"], - next_departure=self._opendata.connections[1]["departure"], - next_on_departure=self._opendata.connections[2]["departure"], + departure=departure_time, + next_departure=next_departure_time.isoformat() + if next_departure_time is not None + else None, + next_on_departure=next_on_departure_time.isoformat() + if next_on_departure_time is not None + else None, train_number=self._opendata.connections[0]["number"], platform=self._opendata.connections[0]["platform"], transfers=self._opendata.connections[0]["transfers"], diff --git a/homeassistant/components/swiss_public_transport/sensor.py b/homeassistant/components/swiss_public_transport/sensor.py index 5d4a6813d2d90f..64487383181330 100644 --- a/homeassistant/components/swiss_public_transport/sensor.py +++ b/homeassistant/components/swiss_public_transport/sensor.py @@ -1,14 +1,18 @@ """Support for transport.opendata.ch.""" from __future__ import annotations -from datetime import timedelta +from datetime import datetime, timedelta import logging from typing import TYPE_CHECKING import voluptuous as vol from homeassistant import config_entries, core -from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity +from homeassistant.components.sensor import ( + PLATFORM_SCHEMA, + SensorDeviceClass, + SensorEntity, +) from homeassistant.config_entries import SOURCE_IMPORT from homeassistant.const import CONF_NAME from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant, callback @@ -107,6 +111,7 @@ class SwissPublicTransportSensor( _attr_icon = "mdi:bus" _attr_has_entity_name = True _attr_translation_key = "departure" + _attr_device_class = SensorDeviceClass.TIMESTAMP def __init__( self, @@ -133,6 +138,6 @@ def _handle_coordinator_update(self) -> None: return super()._handle_coordinator_update() @property - def native_value(self) -> str: + def native_value(self) -> datetime | None: """Return the state of the sensor.""" return self.coordinator.data["departure"]