Merge pull request #5 from droso-hass/new_attributes

New attributes
droso-hass · Nov 14, 2022 · b4794eb · b4794eb
2 parents edb40ee + 8ee465b
commit b4794eb
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 2 deletions.
diff --git a/cli.py b/cli.py
@@ -29,7 +29,7 @@ async def main():
     print("\n")
     print("Traffic:")
     for i in await idfm.get_traffic(stop.id, destination_name=dir):
-        print(f"Line {i.line_id} {i.note} - Destination {i.destination_name}: {i.schedule}") 
+        print(f"Line {i.line_id} {i.note} - Destination {i.destination_name}: {i.schedule} - Currently at stop: {i.at_stop} - Platform: {i.platform} - Status: {i.status}") 
 
     print("\n")
     print("Informations:")

diff --git a/idfm_api/models.py b/idfm_api/models.py
@@ -13,6 +13,21 @@ class TransportType(str, Enum):
     TRAIN = "rail"
     BUS = "bus"
 
+@unique
+class TransportStatus(str, Enum):
+    """
+    Represents the status of a transport
+    """
+    ON_TIME = "onTime"
+    MISSED = "missed"
+    ARRIVED = "arrived"
+    NOT_EXPECTED = "notExpected"
+    DELAYED = "delayed"
+    EARLY = "early"
+    CANCELLED = "cancelled"
+    NO_REPORT = "noReport"
+    UNKNOWN = "unknown"
+
 @dataclass(frozen=True)
 class LineData:
     """
@@ -86,6 +101,9 @@ class TrafficData:
     direction: str
     schedule: datetime
     retarted: bool
+    at_stop: bool
+    platform: str
+    status: str
 
     @staticmethod
     def from_json(data: dict):
@@ -107,14 +125,34 @@ def from_json(data: dict):
         else:
             return None
 
+        try:
+            atstop = data["MonitoredVehicleJourney"]["MonitoredCall"]["VehicleAtStop"]
+        except KeyError:
+            atstop = None
+
+        try:
+            plat = data["MonitoredVehicleJourney"]["MonitoredCall"]["ArrivalPlatformName"]["value"]
+        except KeyError:
+            plat = ""
+
+        if "ArrivalStatus" in data["MonitoredVehicleJourney"]["MonitoredCall"]:
+            status = TransportStatus(data["MonitoredVehicleJourney"]["MonitoredCall"]["ArrivalStatus"])
+        elif "DepartureStatus" in data["MonitoredVehicleJourney"]["MonitoredCall"]:
+            status = TransportStatus(data["MonitoredVehicleJourney"]["MonitoredCall"]["DepartureStatus"])
+        else:
+            status = TransportStatus.UNKNOWN
+
         return TrafficData(
             line_id=data["MonitoredVehicleJourney"]["LineRef"]["value"],
             note=note,
             destination_name=data["MonitoredVehicleJourney"]["DestinationName"][0]["value"],
             destination_id=data["MonitoredVehicleJourney"]["DestinationRef"]["value"],
             direction=dir,
             schedule=sch,
-            retarted=data["MonitoredVehicleJourney"]["MonitoredCall"].get("ArrivalStatus") in [None, "onTime"]
+            retarted=status not in [TransportStatus.ON_TIME, TransportStatus.ARRIVED, TransportStatus.UNKNOWN],
+            at_stop=atstop,
+            platform=plat,
+            status=status
         )
 
     def __eq__(self, other):