From 2e87d0f120cd996afe3a0ed4ce73bf1baff8fc7e Mon Sep 17 00:00:00 2001 From: Edward Betts Date: Wed, 19 Jun 2024 22:27:56 +0100 Subject: [PATCH] Consider railway stations for trip countries --- agenda/types.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) 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