264 lines
8.3 KiB
Python
264 lines
8.3 KiB
Python
# 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.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):
|
|
"""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")
|