Set proper sensor device class for swiss_public_transport

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"],

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"]