diff --git a/agenda/trip.py b/agenda/trip.py index e730ad4..7ad10d2 100644 --- a/agenda/trip.py +++ b/agenda/trip.py @@ -225,12 +225,15 @@ def add_coordinates_for_unbooked_flights( ) -def collect_trip_coordinates(trip: Trip) -> list[StrDict]: - """Extract and de-duplicate airport and station coordinates from trip.""" - stations = {} +def get_locations(trip: Trip) -> dict[str, StrDict]: + """Collect locations of all travel locations in trip.""" + locations: dict[str, StrDict] = { + "station": {}, + "airport": {}, + "ferry_terminal": {}, + } + station_list = [] - airports = {} - ferry_terminals = {} for t in trip.travel: if t["type"] == "train": station_list += [t["from_station"], t["to_station"]] @@ -240,18 +243,33 @@ def collect_trip_coordinates(trip: Trip) -> list[StrDict]: elif t["type"] == "flight": for field in "from_airport", "to_airport": if field in t: - airports[t[field]["iata"]] = t[field] + locations["airport"][t[field]["iata"]] = t[field] else: assert t["type"] == "ferry" for field in "from_terminal", "to_terminal": terminal = t[field] - ferry_terminals[terminal["name"]] = terminal + locations["ferry_terminal"][terminal["name"]] = terminal for s in station_list: - if s["uic"] in stations: + if s["uic"] in locations["station"]: continue - stations[s["uic"]] = s + locations["station"][s["uic"]] = s + return locations + + +def coordinate_dict(item: StrDict, coord_type: str) -> StrDict: + """Build coodinate dict for item.""" + return { + "name": item["name"], + "type": coord_type, + "latitude": item["latitude"], + "longitude": item["longitude"], + } + + +def collect_trip_coordinates(trip: Trip) -> list[StrDict]: + """Extract and de-duplicate travel location coordinates from trip.""" coords = [] src = [ @@ -261,31 +279,14 @@ def collect_trip_coordinates(trip: Trip) -> list[StrDict]: ] for coord_type, item_list in src: coords += [ - { - "name": item["name"], - "type": coord_type, - "latitude": item["latitude"], - "longitude": item["longitude"], - } + coordinate_dict(item, coord_type) for item in item_list if "latitude" in item and "longitude" in item ] - locations = [ - ("station", stations), - ("airport", airports), - ("ferry_terminal", ferry_terminals), - ] - for coord_type, coord_dict in locations: - coords += [ - { - "name": s["name"], - "type": coord_type, - "latitude": s["latitude"], - "longitude": s["longitude"], - } - for s in coord_dict.values() - ] + locations = get_locations(trip) + for coord_type, coord_dict in locations.items(): + coords += [coordinate_dict(s, coord_type) for s in coord_dict.values()] return coords