agenda/trip.py

@@ -225,51 +225,33 @@ def add_coordinates_for_unbooked_flights(
     )
 
 
-def get_locations(trip: Trip) -> dict[str, StrDict]:
+def collect_trip_coordinates(trip: Trip) -> list[StrDict]:
-    """Collect locations of all travel locations in trip."""
+    """Extract and de-duplicate airport and station coordinates from trip."""
-    locations: dict[str, StrDict] = {
+    stations = {}
-        "station": {},
-        "airport": {},
-        "ferry_terminal": {},
-    }
-
     station_list = []
+    airports = {}
+    ferry_terminals = {}
     for t in trip.travel:
-        match t["type"]:
+        if t["type"] == "train":
-            case "train":
+            station_list += [t["from_station"], t["to_station"]]
-                station_list += [t["from_station"], t["to_station"]]
+            for leg in t["legs"]:
-                for leg in t["legs"]:
+                station_list.append(leg["from_station"])
-                    station_list.append(leg["from_station"])
+                station_list.append(leg["to_station"])
-                    station_list.append(leg["to_station"])
+        elif t["type"] == "flight":
-            case "flight":
+            for field in "from_airport", "to_airport":
-                for field in "from_airport", "to_airport":
+                if field in t:
-                    if field in t:
+                    airports[t[field]["iata"]] = t[field]
-                        locations["airport"][t[field]["iata"]] = t[field]
+        else:
-            case "ferry":
+            assert t["type"] == "ferry"
-                for field in "from_terminal", "to_terminal":
+            for field in "from_terminal", "to_terminal":
-                    terminal = t[field]
+                terminal = t[field]
-                    locations["ferry_terminal"][terminal["name"]] = terminal
+                ferry_terminals[terminal["name"]] = terminal
 
     for s in station_list:
-        if s["uic"] in locations["station"]:
+        if s["uic"] in stations:
             continue
-        locations["station"][s["uic"]] = s
+        stations[s["uic"]] = s
 
-    return locations
-
-
-def coordinate_dict(item: StrDict, coord_type: str) -> StrDict:
-    """Build coodinate dict for item."""
-    return {
-        "name": item["name"],
-        "type": coord_type,
-        "latitude": item["latitude"],
-        "longitude": item["longitude"],
-    }
-
-
-def collect_trip_coordinates(trip: Trip) -> list[StrDict]:
-    """Extract and de-duplicate travel location coordinates from trip."""
     coords = []
 
     src = [
@@ -279,14 +261,31 @@ def collect_trip_coordinates(trip: Trip) -> list[StrDict]:
     ]
     for coord_type, item_list in src:
         coords += [
-            coordinate_dict(item, coord_type)
+            {
+                "name": item["name"],
+                "type": coord_type,
+                "latitude": item["latitude"],
+                "longitude": item["longitude"],
+            }
             for item in item_list
             if "latitude" in item and "longitude" in item
         ]
 
-    locations = get_locations(trip)
+    locations = [
-    for coord_type, coord_dict in locations.items():
+        ("station", stations),
-        coords += [coordinate_dict(s, coord_type) for s in coord_dict.values()]
+        ("airport", airports),
+        ("ferry_terminal", ferry_terminals),
+    ]
+    for coord_type, coord_dict in locations:
+        coords += [
+            {
+                "name": s["name"],
+                "type": coord_type,
+                "latitude": s["latitude"],
+                "longitude": s["longitude"],
+            }
+            for s in coord_dict.values()
+        ]
 
     return coords