Add UnknownStation Exception

This commit is contained in:
Edward Betts 2024-03-27 16:34:24 +00:00
parent d690442f0f
commit 826eafbc86

View file

@ -9,6 +9,12 @@ from agenda import travel
from agenda.types import StrDict, Trip from agenda.types import StrDict, Trip
class UnknownStation(Exception):
"""Unknown station."""
pass
def load_travel(travel_type: str, data_dir: str) -> list[StrDict]: def load_travel(travel_type: str, data_dir: str) -> list[StrDict]:
"""Read flight and train journeys.""" """Read flight and train journeys."""
items = travel.parse_yaml(travel_type + "s", data_dir) items = travel.parse_yaml(travel_type + "s", data_dir)
@ -24,8 +30,10 @@ def load_trains(data_dir: str) -> list[StrDict]:
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:
assert train["from"] in by_name if train["from"] not in by_name:
assert train["to"] in 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["from_station"] = by_name[train["from"]]
train["to_station"] = by_name[train["to"]] train["to_station"] = by_name[train["to"]]