diff --git a/captaintrain.py b/captaintrain.py index b92b68b..34c2119 100644 --- a/captaintrain.py +++ b/captaintrain.py @@ -15,6 +15,29 @@ CT_HEADERS = { } +# Structures + +class Station: + + def __init__(self, id, is_sellable, name, slug, score, info, latitude, longitude): + self._id = id + self._is_sellable = is_sellable + self._name = name + self._slug = slug + self._score = score + self._info = info + self._latitude = latitude + self._longitude = longitude + + @property + def id(self): + return self._id + + @property + def name(self): + return self._name + + # API calls def stations(name): @@ -27,8 +50,8 @@ def stations(name): with urllib.request.urlopen(req) as res: stations = json.loads(res.read().decode())["stations"] sorted(stations, key=lambda station: station["score"], reverse=True) - return stations - return [] + for station in stations: + yield Station(**station) def station(name): @@ -64,9 +87,9 @@ def search(departure, arrival, departure_date, return_date=None): ], "systems": ["sncf","db","busbud","idtgv","ouigo","trenitalia","ntv","hkx","renfe","timetable"], "exchangeable_part": None, - "departure_station_id": departure["id"], + "departure_station_id": departure.id, "via_station_id": None, - "arrival_station_id": arrival["id"], + "arrival_station_id": arrival.id, "exchangeable_pnr_id": None } }).encode()) diff --git a/cheapest_ticket.py b/cheapest_ticket.py index 397cb9e..ef55903 100755 --- a/cheapest_ticket.py +++ b/cheapest_ticket.py @@ -20,8 +20,8 @@ if __name__ == "__main__": else: departure_time = ct.parse_datetime(departure_time) - print("From:", departure["name"], departure["id"]) - print("To:", arrival["name"], arrival["id"]) + print("From:", departure.name, departure.id) + print("To:", arrival.name, arrival.id) print("Departure:", departure_time) #with open("res.json") as f: