Use pattern matching: train/flight/ferry

This commit is contained in:
Edward Betts 2024-10-02 15:36:26 +01:00
parent 7d803e0267
commit 868c1407b5

View file

@ -235,17 +235,17 @@ def get_locations(trip: Trip) -> dict[str, StrDict]:
station_list = [] station_list = []
for t in trip.travel: for t in trip.travel:
if t["type"] == "train": match t["type"]:
case "train":
station_list += [t["from_station"], t["to_station"]] station_list += [t["from_station"], t["to_station"]]
for leg in t["legs"]: for leg in t["legs"]:
station_list.append(leg["from_station"]) station_list.append(leg["from_station"])
station_list.append(leg["to_station"]) station_list.append(leg["to_station"])
elif t["type"] == "flight": case "flight":
for field in "from_airport", "to_airport": for field in "from_airport", "to_airport":
if field in t: if field in t:
locations["airport"][t[field]["iata"]] = t[field] locations["airport"][t[field]["iata"]] = t[field]
else: case "ferry":
assert t["type"] == "ferry"
for field in "from_terminal", "to_terminal": for field in "from_terminal", "to_terminal":
terminal = t[field] terminal = t[field]
locations["ferry_terminal"][terminal["name"]] = terminal locations["ferry_terminal"][terminal["name"]] = terminal