From f423fcdcbe28c84a9808a0877383e599f134e08d Mon Sep 17 00:00:00 2001
From: Edward Betts <edward@4angle.com>
Date: Sat, 3 Aug 2024 14:49:21 +0800
Subject: [PATCH] Split up train and flight loading

Reduce complexity of train and flight loading functions by splitting
code out into separate functions.
---
 agenda/trip.py | 44 +++++++++++++++++++++++++++++---------------
 1 file changed, 29 insertions(+), 15 deletions(-)

diff --git a/agenda/trip.py b/agenda/trip.py
index 850da77..616e29e 100644
--- a/agenda/trip.py
+++ b/agenda/trip.py
@@ -28,6 +28,19 @@ def load_travel(travel_type: str, plural: str, data_dir: str) -> list[StrDict]:
     return items
 
 
+def process_train_leg(
+    leg: StrDict,
+    by_name: StrDict,
+    route_distances: travel.RouteDistances | None,
+) -> None:
+    """Process train leg."""
+    assert leg["from"] in by_name and leg["to"] in by_name
+    leg["from_station"], leg["to_station"] = by_name[leg["from"]], by_name[leg["to"]]
+
+    if route_distances:
+        travel.add_leg_route_distance(leg, route_distances)
+
+
 def load_trains(
     data_dir: str, route_distances: travel.RouteDistances | None = None
 ) -> list[StrDict]:
@@ -45,13 +58,7 @@ def load_trains(
         train["to_station"] = by_name[train["to"]]
 
         for leg in train["legs"]:
-            assert leg["from"] in by_name
-            assert leg["to"] in by_name
-            leg["from_station"] = by_name[leg["from"]]
-            leg["to_station"] = by_name[leg["to"]]
-
-            if route_distances:
-                travel.add_leg_route_distance(leg, route_distances)
+            process_train_leg(leg, by_name=by_name, route_distances=route_distances)
 
         if all("distance" in leg for leg in train["legs"]):
             train["distance"] = sum(leg["distance"] for leg in train["legs"])
@@ -90,6 +97,20 @@ def depart_datetime(item: StrDict) -> datetime:
     return datetime.combine(depart, time.min).replace(tzinfo=ZoneInfo("UTC"))
 
 
+def process_flight(
+    flight: StrDict, iata: dict[str, str], airports: list[StrDict]
+) -> None:
+    """Add airport detail, airline name and distance to flight."""
+    if flight["from"] in airports:
+        flight["from_airport"] = airports[flight["from"]]
+    if flight["to"] in airports:
+        flight["to_airport"] = airports[flight["to"]]
+    if "airline" in flight:
+        flight["airline_name"] = iata.get(flight["airline"], "[unknown]")
+
+    flight["distance"] = travel.flight_distance(flight)
+
+
 def load_flight_bookings(data_dir: str) -> list[StrDict]:
     """Load flight bookings."""
     bookings = load_travel("flight", "flights", data_dir)
@@ -98,14 +119,7 @@ def load_flight_bookings(data_dir: str) -> list[StrDict]:
     airports = travel.parse_yaml("airports", data_dir)
     for booking in bookings:
         for flight in booking["flights"]:
-            if flight["from"] in airports:
-                flight["from_airport"] = airports[flight["from"]]
-            if flight["to"] in airports:
-                flight["to_airport"] = airports[flight["to"]]
-            if "airline" in flight:
-                flight["airline_name"] = iata.get(flight["airline"], "[unknown]")
-
-            flight["distance"] = travel.flight_distance(flight)
+            process_flight(flight, iata, airports)
     return bookings