Split travel YAML parsing into function.

agenda/travel.py

@@ -1,6 +1,7 @@
 """Travel."""
 
 import os
+import typing
 
 import yaml
 
@@ -8,10 +9,17 @@ from .types import Event
 
 Leg = dict[str, str]
 
+TravelList = list[dict[str, typing.Any]]
+
+
+def parse_yaml(travel_type: str, data_dir: str) -> TravelList:
+    """Parse flights YAML and return list of travel."""
+    filepath = os.path.join(data_dir, travel_type + ".yaml")
+    return typing.cast(TravelList, yaml.safe_load(open(filepath)))
+
 
 def get_flights(data_dir: str) -> list[Event]:
     """Get travel events."""
-    filepath = os.path.join(data_dir, "flights.yaml")
     return [
         Event(
             date=item["depart"],
@@ -20,7 +28,7 @@ def get_flights(data_dir: str) -> list[Event]:
             title=f'flight from {item["from"]} to {item["to"]} ({flight_number(item)})',
             url=item.get("url"),
         )
-        for item in yaml.safe_load(open(filepath))
+        for item in parse_yaml("flights", data_dir)
         if item["depart"].date()
     ]
 
@@ -28,8 +36,7 @@ def get_flights(data_dir: str) -> list[Event]:
 def get_trains(data_dir: str) -> list[Event]:
     """Get train events."""
     events: list[Event] = []
-    filepath = os.path.join(data_dir, "trains.yaml")
+    for item in parse_yaml("trains", data_dir):
-    for item in yaml.safe_load(open(filepath)):
         events += [
             Event(
                 date=leg["depart"],