Working with stub classes

captaintrain.py

@@ -19,15 +19,9 @@ CT_HEADERS = {
 
 class Station:
 
-    def __init__(self, id, is_sellable, name, slug, score, info, latitude, longitude):
+    def __init__(self, id, name, slug, info, latitude=0, longitude=0, score=1, is_sellable=True, parent_name=None, address=None):
         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):
@@ -38,6 +32,156 @@ class Station:
         return self._name
 
 
+class ComfortClass:
+
+    def __init__(self, id, name, cents, currency, condition_id, segment_id, default=False, is_available=False, options={}, reservation=None, title="", description=""):
+        pass
+
+
+class Condition:
+
+    def __init__(self, id, name, short_description="", comfort_class=None, segment_id=None):
+        pass
+
+
+class Folder:
+
+    def __init__(self, id, search_id, trip_ids, arrival_date, cents, travel_class, comfort, created_at, has_round_trip_fare, arrival_station_id, system, departure_date, is_sellable, is_birthdate_required, flexibility, digest, direction, currency, departure_station_id, is_only_possible_travel_class=False):
+        pass
+
+
+class Passenger:
+
+    def __init__(self, id, label):
+        pass
+
+
+class Segment:
+
+    def __init__(self, id, trip_id, train_number, digest, options, departure_date, train_name, arrival_station_id, comfort_class_ids, departure_station_id, travel_class, arrival_date, condition_id, carrier, train, co2_emission=0, reservation=False):
+        self._id = id
+        self._carrier = carrier
+        self._train = train
+        self._train_name = train_name
+        self._train_number = train_number
+
+    def __str__(self):
+        return "%s %s %s" % (self.carrier, self.train_name, self.train_number)
+
+    @property
+    def id(self):
+        return self._id
+
+    @property
+    def carrier(self):
+        return self._carrier
+
+    @property
+    def train(self):
+        return self._train
+
+    @property
+    def train_name(self):
+        return self._train_name
+
+    @property
+    def train_number(self):
+        return self._train_number
+
+
+class Trip:
+
+    def __init__(self, get_station, get_segment, id, departure_date, arrival_date, cents, currency, folder_id, segment_ids, digest, departure_station_id, arrival_station_id, passenger_id):
+        self._price = cents / 100
+        self._currency = currency
+        self._digest = digest
+        self._departure_date = parse_datetime(departure_date)
+        self._arrival_date = parse_datetime(arrival_date)
+        self._departure_station = get_station(departure_station_id)
+        self._arrival_station = get_station(arrival_station_id)
+        self._segments = [get_segment(segment_id) for segment_id in segment_ids]
+
+
+    def __str__(self):
+        return "From %s at %s to %s at %s // duration: %s // price: %s %s // %s" % (
+            self.departure_station.name, self.departure_date,
+            self.arrival_station.name, self.arrival_date,
+            self.arrival_date - self.departure_date,
+            self.price, self.currency,
+            " + ".join([str(s) for s in self.segments])
+            )
+
+
+    @property
+    def price(self):
+        return self._price
+
+    @property
+    def currency(self):
+        return self._currency
+
+    @property
+    def digest(self):
+        return self._digest
+
+    @property
+    def departure_date(self):
+        return self._departure_date
+
+    @property
+    def arrival_date(self):
+        return self._arrival_date
+
+    @property
+    def departure_station(self):
+        return self._departure_station
+
+    @property
+    def arrival_station(self):
+        return self._arrival_station
+
+    @property
+    def segments(self):
+        return self._segments
+
+
+class Search:
+
+    def __init__(self, comfort_classes, conditions, folders, passengers, search, segments, stations, trips):
+        self._comfort_classes = [ComfortClass(**comfort_class) for comfort_class in comfort_classes]
+        self._conditions = [Condition(**condition) for condition in conditions]
+        self._folders = [Folder(**folder) for folder in folders]
+        self._passengers = [Passenger(**passenger) for passenger in passengers]
+        self._search = search
+        self._segments = [Segment(**segment) for segment in segments]
+        self._stations = [Station(**station) for station in stations]
+        self._trips = [Trip(self.get_station, self.get_segment, **trip) for trip in trips]
+
+
+    @property
+    def trips(self):
+        return self._trips
+
+    @property
+    def stations(self):
+        return self._stations
+
+    @property
+    def segments(self):
+        return self._segments
+
+
+    def get_station(self, station_id):
+        for station in self._stations:
+            if station.id == station_id:
+                return station
+
+    def get_segment(self, segment_id):
+        for segment in self._segments:
+            if segment.id == segment_id:
+                return segment
+
+
 # API calls
 
 def stations(name):
@@ -94,7 +238,7 @@ def search(departure, arrival, departure_date, return_date=None):
                                      }
                                  }).encode())
     with urllib.request.urlopen(req) as res:
-        return json.loads(res.read().decode())
+        return Search(**json.loads(res.read().decode()))
     return None
 
 
@@ -104,48 +248,16 @@ def cheapest_trips(trips):
     min_trips = []
 
     for trip in trips:
-        if len(min_trips) == 0 or min_trips[0]["cents"] > trip["cents"]:
+        if len(min_trips) == 0 or min_trips[0].price > trip.price:
             min_trips = [trip]
-        elif min_trips[0]["cents"] == trip["cents"] and trip["digest"] not in map(lambda x: x["digest"], min_trips):
+        elif min_trips[0].price == trip.price and trip.digest not in map(lambda x: x.digest, min_trips):
             min_trips.append(trip)
 
     return min_trips
 
 
-def display_segment(segment, stations):
-    def get_station(station_id):
-        for station in stations:
-            if station["id"] == station_id:
-                return station
-        return None
-
-
 def parse_datetime(s):
     try:
         return datetime.strptime(s, "%Y-%m-%dT%H:%M:%S%z")
     except ValueError:
         return datetime.strptime(s[::-1].replace(":", "", 1)[::-1], "%Y-%m-%dT%H:%M:%S%z")
-
-
-def display_trip(trip, stations, segments):
-    def get_station(station_id):
-        for station in stations:
-            if station["id"] == station_id:
-                return station
-        return None
-
-    def get_segments():
-        for segment in segments:
-            if segment["trip_id"] == trip["id"]:
-                yield segment
-
-    departure_date = parse_datetime(trip["departure_date"])
-    arrival_date = parse_datetime(trip["arrival_date"])
-
-    return "From %s at %s to %s at %s // duration: %s // price: %s %s // %s" % (
-        get_station(trip["departure_station_id"])["name"], departure_date,
-        get_station(trip["arrival_station_id"])["name"], arrival_date,
-        arrival_date - departure_date,
-        trip["cents"]/100, trip["currency"],
-        " + ".join(map(lambda segment: "%s %s %s" % (segment["carrier"], segment["train_name"], segment["train_number"]), get_segments()))
-    )

cheapest_ticket.py

@@ -26,11 +26,11 @@ if __name__ == "__main__":
 
         #with open("res.json") as f:
         #    import json
-        #    res = json.load(f)
+        #    res = ct.Search(**json.load(f))
 
         res = ct.search(departure, arrival, departure_time)
-        min_trips = ct.cheapest_trips(res["trips"])
+        min_trips = ct.cheapest_trips(res.trips)
         for trip in min_trips:
-            print(ct.display_trip(trip, res["stations"], res["segments"]))
+            print(trip)
     else:
         print("usage")