diff --git a/agenda/__init__.py b/agenda/__init__.py index c07f058..79a9589 100644 --- a/agenda/__init__.py +++ b/agenda/__init__.py @@ -360,13 +360,28 @@ def get_accommodation(from_date: date, filepath: str) -> list[Event]: return events -def get_travel(from_date: date, method: str, filepath: str) -> list[Event]: +Leg = dict[str, str] + + +def get_travel( + from_date: date, + method: str, + filepath: str, + extra: typing.Callable[[Leg], str] | None = None, +) -> list[Event]: """Get travel events.""" + + def title(item: Leg) -> str: + ret = f'{method} from {item["from"]} to {item["to"]}' + if extra: + ret += f" ({extra(item)})" + return ret + return [ Event( date=item["depart"], name="transport", - title=f'{method} from {item["from"]} to {item["to"]}', + title=title(item), url=item.get("url"), ) for item in yaml.safe_load(open(filepath)) @@ -374,11 +389,22 @@ def get_travel(from_date: date, method: str, filepath: str) -> list[Event]: ] +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 get_all_travel_events(from_date: date) -> list[Event]: """Get all flights and rail journeys.""" data_dir = config["data"]["personal-data"] trains = get_travel(from_date, "train", os.path.join(data_dir, "trains.yaml")) - flights = get_travel(from_date, "flight", os.path.join(data_dir, "flights.yaml")) + flights = get_travel( + from_date, "flight", os.path.join(data_dir, "flights.yaml"), extra=flight_number + ) return trains + flights