Introducing Station class

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())

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: