Merge pull request #28 from M-See/main

Update sensor.py
FaserF · Dec 28, 2023 · e438bfe · e438bfe
2 parents 2629bd5 + ad15a4b
commit e438bfe
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/custom_components/deutschebahn/sensor.py b/custom_components/deutschebahn/sensor.py
@@ -103,14 +103,12 @@ def extra_state_attributes(self):
         """Return the state attributes."""
         if len(self.connections) > 0:
             connections = self.connections[0].copy()
-            if len(self.connections) > 1:
-                connections["next"] = self.connections[1]["departure"]
-                connections["next_delay"] = self.connections[1]["delay"]
-                connections["next_canceled"] = self.connections[1]["canceled"]
-            if len(self.connections) > 2:
-                connections["next_on"] = self.connections[2]["departure"]
-                connections["next_on_delay"] = self.connections[2]["delay"]
-                connections["next_on_canceled"] = self.connections[2]["canceled"]
+            for cons in range(1,len(self.connections)):
+                if len(self.connections) > cons:
+                    connections[f"next_{cons}"] = self.connections[cons]["departure"]
+                    connections[f"next_{cons}_delay"] = self.connections[cons]["delay"]
+                    connections[f"next_{cons}_canceled"] = self.connections[cons]["canceled"]
+                    connections[f"next_{cons}_arrival"] = self.connections[cons]["arrival"]
         else:
             connections = None
         connections["departures"] = self.connections