display new attributes within every connection #25

FaserF · Jun 9, 2024 · 9f2d6a7 · 9f2d6a7
1 parent d790650
commit 9f2d6a7
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/custom_components/deutschebahn/sensor.py b/custom_components/deutschebahn/sensor.py
@@ -98,15 +98,24 @@ def native_value(self):
     @property
     def extra_state_attributes(self):
         """Return the state attributes."""
-        attributes = {
-            "departures": self.connections,
-        }
+        attributes = {}
         if self.connections:
-            first_connection = self.connections[0]
-            departure_current = first_connection.get("departure") + " +" + str(first_connection.get("delay", 0))
-            arrival_current = first_connection.get("arrival") + " +" + str(first_connection.get("delay_arrival", 0))
-            attributes["departure_current"] = departure_current
-            attributes["arrival_current"] = arrival_current
+            for con in self.connections:
+                if "departure" in con and "arrival" in con:
+                    # Parse departure and arrival times
+                    departure_time = dt_util.parse_time(con.get("departure"))
+                    arrival_time = dt_util.parse_time(con.get("arrival"))
+                    if departure_time and arrival_time:
+                        # Create datetime objects for departure and arrival times
+                        departure_datetime = datetime.combine(datetime.now().date(), departure_time)
+                        arrival_datetime = datetime.combine(datetime.now().date(), arrival_time)
+                        # Apply delays
+                        corrected_departure_time = departure_datetime + timedelta(minutes=con.get("delay", 0))
+                        corrected_arrival_time = arrival_datetime + timedelta(minutes=con.get("delay_arrival", 0))
+                        # Format and add current departure and arrival times
+                        con["departure_current"] = corrected_departure_time.strftime("%H:%M")
+                        con["arrival_current"] = corrected_arrival_time.strftime("%H:%M")
+        attributes["departures"] = self.connections
         return attributes
 
     async def async_update(self):