diff --git a/agenda/types.py b/agenda/types.py index b37f5ba..92ea0a5 100644 --- a/agenda/types.py +++ b/agenda/types.py @@ -160,13 +160,18 @@ class Trip: items.append(country) for item in self.travel: - if item["type"] != "flight": - continue + travel_countries = set() + if item["type"] == "flight": + for key in "from_airport", "to_airport": + c = item[key]["country"] + travel_countries.add(c) + if item["type"] == "train": + for leg in item["legs"]: + for key in "from_station", "to_station": + c = leg[key]["country"] + travel_countries.add(c) - for key in "from_airport", "to_airport": - c = item[key]["country"] - if c in seen: - continue + for c in travel_countries - seen: seen.add(c) country = agenda.get_country(c) assert country