Initial commit

captaintrain.py

@@ -0,0 +1,128 @@
+# 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",
+}
+
+
+# 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)
+        return stations
+    return []
+
+
+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()))
+    )