Add train route distance info

agenda/travel.py

@@ -1,12 +1,14 @@
 """Travel."""
 
+import json
 import os
 import typing
 
 import flask
 import yaml
+from geopy.distance import geodesic
 
-from .types import Event
+from .types import Event, StrDict
 
 Leg = dict[str, str]
 
@@ -63,3 +65,47 @@ def flight_number(flight: Leg) -> str:
 def all_events(data_dir: str) -> list[Event]:
     """Get all flights and rail journeys."""
     return get_trains(data_dir) + get_flights(data_dir)
+
+
+RouteDistances = dict[tuple[str, str], float]
+
+
+def train_leg_distance(geojson_data: StrDict) -> float:
+    """Calculate the total length of a LineString in kilometers from GeoJSON data."""
+    # Extract coordinates
+    first_object = geojson_data["features"][0]["geometry"]
+    assert first_object["type"] in ("LineString", "MultiLineString")
+
+    if first_object["type"] == "LineString":
+        coord_list = [first_object["coordinates"]]
+    else:
+        first_object["type"] == "MultiLineString"
+        coord_list = first_object["coordinates"]
+    # pprint(coordinates)
+
+    total_length_km = 0.0
+
+    for coordinates in coord_list:
+        total_length_km += sum(
+            float(geodesic(coordinates[i], coordinates[i + 1]).km)
+            for i in range(len(coordinates) - 1)
+        )
+
+    return total_length_km
+
+
+def load_route_distances(data_dir: str) -> RouteDistances:
+    """Load cache of route distances."""
+    route_distances: RouteDistances = {}
+    with open(os.path.join(data_dir, "route_distances.json")) as f:
+        for s1, s2, dist in json.load(f):
+            route_distances[(s1, s2)] = dist
+
+    return route_distances
+
+
+def add_leg_route_distance(leg: StrDict, route_distances: RouteDistances) -> None:
+    s1, s2 = sorted([leg["from"], leg["to"]])
+    dist = route_distances.get((s1, s2))
+    if dist:
+        leg["distance"] = dist