ct-tarifs/captaintrain.py

# 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")