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