Catch missing station in train leg

Raise UnknownStation for missing stations in train journey leg.
This commit is contained in:
Edward Betts 2024-09-25 12:15:13 +01:00
parent 488ecf8b71
commit 96ec2b7d89

View file

@ -28,17 +28,18 @@ def load_travel(travel_type: str, plural: str, data_dir: str) -> list[StrDict]:
return items return items
def process_train_leg( def get_station(name: str, by_name: dict[str, StrDict]) -> StrDict:
leg: StrDict, """Get station by name."""
by_name: StrDict, try:
route_distances: travel.RouteDistances | None, return by_name[name]
) -> None: except IndexError:
"""Process train leg.""" raise UnknownStation(name)
assert leg["from"] in by_name and leg["to"] in by_name
leg["from_station"], leg["to_station"] = by_name[leg["from"]], by_name[leg["to"]]
if route_distances:
travel.add_leg_route_distance(leg, route_distances) def add_station_objects(item: StrDict, by_name: dict[str, StrDict]) -> None:
"""Lookup stations and add to train or leg."""
item["from_station"] = get_station(item["from"], by_name)
item["to_station"] = get_station(item["to"], by_name)
def load_trains( def load_trains(
@ -50,15 +51,11 @@ def load_trains(
by_name = {station["name"]: station for station in stations} by_name = {station["name"]: station for station in stations}
for train in trains: for train in trains:
if train["from"] not in by_name: add_station_objects(train, by_name)
raise UnknownStation(train["from"])
if train["to"] not in by_name:
raise UnknownStation(train["to"])
train["from_station"] = by_name[train["from"]]
train["to_station"] = by_name[train["to"]]
for leg in train["legs"]: for leg in train["legs"]:
process_train_leg(leg, by_name=by_name, route_distances=route_distances) add_station_objects(leg, by_name)
if route_distances:
travel.add_leg_route_distance(leg, route_distances)
if all("distance" in leg for leg in train["legs"]): if all("distance" in leg for leg in train["legs"]):
train["distance"] = sum(leg["distance"] for leg in train["legs"]) train["distance"] = sum(leg["distance"] for leg in train["legs"])