agenda/agenda/travel.py

65 lines
1.7 KiB
Python

"""Travel."""
import os
import typing
import yaml
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."""
return [
Event(
date=item["depart"],
end_date=item.get("arrive"),
name="transport",
title=f'✈️ {item["from"]} to {item["to"]} ({flight_number(item)})',
url=item.get("url"),
)
for item in parse_yaml("flights", data_dir)
if item["depart"].date()
]
def get_trains(data_dir: str) -> list[Event]:
"""Get train events."""
events: list[Event] = []
for item in parse_yaml("trains", data_dir):
events += [
Event(
date=leg["depart"],
end_date=leg["arrive"],
name="transport",
title=f'🚆 {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"]
# make sure this is the airline code, not the airline name
assert " " not in airline_code and not any(c.islower() for c in airline_code)
return airline_code + flight["flight_number"]
def all_events(data_dir: str) -> list[Event]:
"""Get all flights and rail journeys."""
return get_trains(data_dir) + get_flights(data_dir)