# 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, 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): """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 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]["cents"] > trip["cents"]: min_trips = [trip] elif min_trips[0]["cents"] == trip["cents"] 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())) )