diff --git a/agenda/travel.py b/agenda/travel.py index f347801..29beff4 100644 --- a/agenda/travel.py +++ b/agenda/travel.py @@ -38,6 +38,25 @@ def get( ] +def get_trains(from_date: date, filepath: str) -> list[Event]: + """Get train events.""" + events: list[Event] = [] + for item in yaml.safe_load(open(filepath)): + if item["depart"].date() < from_date: + continue + events += [ + Event( + date=leg["depart"], + end_date=leg["arrive"], + name="transport", + title=f'train from {leg["from"]} to {leg["to"]}', + url=item.get("url"), + ) + for leg in item["legs"] + ] + return events + + def flight_number(flight: Leg) -> str: """Flight number.""" airline_code = flight["airline"] @@ -49,7 +68,7 @@ def flight_number(flight: Leg) -> str: def all_events(from_date: date, data_dir: str) -> list[Event]: """Get all flights and rail journeys.""" - trains = get(from_date, "train", os.path.join(data_dir, "trains.yaml")) + trains = get_trains(from_date, os.path.join(data_dir, "trains.yaml")) flights = get( from_date, "flight", os.path.join(data_dir, "flights.yaml"), extra=flight_number )