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, 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()))
    )