diff --git a/agenda/travel.py b/agenda/travel.py index 302e2c0..bc1771d 100644 --- a/agenda/travel.py +++ b/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 yaml.safe_load(open(filepath)): + for item in parse_yaml("trains", data_dir): events += [ Event( date=leg["depart"],