diff --git a/captaintrain.py b/captaintrain.py index 1ead58e..aa25870 100644 --- a/captaintrain.py +++ b/captaintrain.py @@ -87,6 +87,7 @@ class Station: 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._name = name + self._parent_name = parent_name @property def id(self): @@ -96,6 +97,10 @@ class Station: def name(self): return self._name + @property + def parent_name(self): + return self._parent_name if self._parent_name is not None else self._name + class ComfortClass: @@ -123,12 +128,16 @@ class Passenger: 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): + def __init__(self, get_station, 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 + self._departure_date = parse_datetime(departure_date) + self._departure_station = get_station(departure_station_id) + self._arrival_date = parse_datetime(arrival_date) + self._arrival_station = get_station(arrival_station_id) def __str__(self): return "%s %s %s" % (self.carrier, self.train_name, self.train_number) @@ -153,6 +162,20 @@ class Segment: def train_number(self): return self._train_number + @property + def departure_date(self): + return self._departure_date + @property + def departure_station(self): + return self._departure_station + + @property + def arrival_date(self): + return self._arrival_date + @property + def arrival_station(self): + return self._arrival_station + class Trip: @@ -168,12 +191,25 @@ class Trip: def __str__(self): - return "From %s at %s to %s at %s // duration: %s // price: %s %s // %s" % ( + via = [] + segmnts = [] + last_seg = None + for seg in self.segments: + if last_seg is not None: + via.append(seg.departure_station.parent_name) + if last_seg.arrival_station.id != seg.departure_station.id: + segmnts.append("-> connection time: %s from %s to %s ->" % (seg.departure_date - last_seg.arrival_date, last_seg.arrival_station.name, seg.departure_station.name)) + else: + segmnts.append("-> connection time: %s in %s ->" % (seg.departure_date - last_seg.arrival_date, seg.departure_station.name)) + segmnts.append(str(seg)) + last_seg = seg + return "From %s at %s to %s at %s%s // duration: %s // price: %s %s // %s" % ( self.departure_station.name, self.departure_date, self.arrival_station.name, self.arrival_date, + " via " + ", ".join(via) if len(via) else "", self.duration, self.price, self.currency, - " + ".join([str(s) for s in self.segments]) + " ".join(segmnts) ) @@ -222,8 +258,8 @@ class Search: 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._segments = [Segment(self.get_station, **segment) for segment in segments] self._trips = [Trip(self.get_station, self.get_segment, **trip) for trip in trips]