# Import stuff from datetime import datetime import json import urllib.request from urllib.parse import quote # Constants CT_HEADERS = { "X-CT-Version": "58380fc3", "X-CT-Timestamp": "1467966546", "X-CT-Client-Id": "cb7a42cd-efe2-42c3-a865-960396510735", } # Structures 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 @property def id(self): return self._id @property def name(self): 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.duration, 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 duration(self): return self._arrival_date - self._departure_date @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): """Get a list of stations matching given name""" headers = {"Accept": "application/json, text/javascript, */*; q=0.01"} headers.update(CT_HEADERS) req = urllib.request.Request('https://www.captaintrain.com/api/v5/stations?context=search&q=' + quote(name), headers=headers) with urllib.request.urlopen(req) as res: stations = json.loads(res.read().decode())["stations"] sorted(stations, key=lambda station: station["score"], reverse=True) for station in stations: yield Station(**station) def station(name): """Get the closest name matching station""" for station in stations(name): return station def search(departure, arrival, departure_date, return_date=None): "Search for a trip" headers = { "Accept": "application/json, text/javascript, */*; q=0.01", "Content-Type": "application/json; charset=utf-8", } headers.update(CT_HEADERS) req = urllib.request.Request('https://www.captaintrain.com/api/v5/search', headers=headers, data=json.dumps({ "search": { "departure_date": departure_date.strftime("%Y-%m-%dT%H:%M:%S%z") if departure_date else None, "return_date": return_date.strftime("%Y-%m-%dT%H:%M:%S%z") if return_date else None, "passengers": [ { "id": "d7c95386-9ed2-4366-8292-610d821940a3", "label": "adult", "age": 27, "cards": [ { "reference": "SNCF.Carte1225" } ], "cui": None } ], "systems": ["sncf","db","busbud","idtgv","ouigo","trenitalia","ntv","hkx","renfe","timetable"], "exchangeable_part": None, "departure_station_id": departure.id, "via_station_id": None, "arrival_station_id": arrival.id, "exchangeable_pnr_id": None } }).encode()) with urllib.request.urlopen(req) as res: return Search(**json.loads(res.read().decode())) return None # Convinience functions def cheapest_trips(trips): min_trips = [] for trip in trips: if len(min_trips) == 0 or min_trips[0].price > trip.price: min_trips = [trip] 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 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")