Use argparse and improve search

captaintrain.py

@@ -14,6 +14,71 @@ CT_HEADERS = {
     "X-CT-Client-Id": "cb7a42cd-efe2-42c3-a865-960396510735",
 }
 
+CT_CARDS = [
+    "Eurostar.FrequentTraveller",
+    "Thalys.TheCard",
+    "Thalys.ThePassBusiness",
+    "Thalys.ThePassPremium",
+    "Thalys.ThePassWeekEnd",
+    "SNCF.Carte1225",
+    "SNCF.CarteEscapades",
+    "SNCF.CarteSenior",
+    "SNCF.CarteEnfantPlus",
+    "SNCF.CarteEnfantFamille",
+    "SNCF.CarteFamilleNombreuse30",
+    "SNCF.CarteFamilleNombreuse40",
+    "SNCF.CarteFamilleNombreuse50",
+    "SNCF.CarteFamilleNombreuse75",
+    "SNCF.CarteMilitaireSecondClass",
+    "SNCF.CarteMilitaireFirstClass",
+    "SNCF.CarteFamilleMilitaire",
+    "SNCF.AbonnementForfaitSecondClass",
+    "SNCF.AbonnementForfaitFirstClass",
+    "SNCF.AbonnementFrequenceSecondClass",
+    "SNCF.AbonnementFrequenceFirstClass",
+    "SNCF.CarteVoyageur",
+    "SNCF.CarteGrandVoyageur",
+    "SNCF.CarteGrandVoyageurPlus",
+    "SNCF.CarteGrandVoyageurLeClub",
+    "DB.BahnCard25SecondClass",
+    "DB.BahnCard25FirstClass",
+    "DB.BahnCard50SecondClass",
+    "DB.BahnCard50FirstClass",
+    "DB.BahnCard25BusinessSecondClass",
+    "DB.BahnCard25BusinessFirstClass",
+    "DB.BahnCard50BusinessSecondClass",
+    "DB.BahnCard50BusinessFirstClass",
+    "DB.BahnBonusCard",
+    "Renfe.CarneJoven",
+    "Renfe.TarjetaDorada",
+    "Renfe.TarjetaMasRenfeJoven50",
+    "Renfe.TarjetaMasRenfe",
+    "Renfe.TarjetaMasRenfePlata",
+    "Renfe.TarjetaMasRenfeOro",
+    "Renfe.TarjetaMasRenfePremium",
+    "Trenitalia.CartaFreccia",
+    "Trenitalia.CartaVerde",
+    "Trenitalia.CartaArgento",
+    "NTV.ItaloPiu",
+    "NS.VoordeelurenabonnementRailPlus",
+    "NS.Voordeelurenabonnement",
+    "CFF.DemiTarifRailPlus",
+    "CFF.DemiTarif",
+    "OBB.VORTEILScard",
+]
+
+CT_SYSTEMS = [
+    "sncf",
+    "db",
+    "busbud",
+    "idtgv",
+    "ouigo",
+    "trenitalia",
+    "ntv",
+    "hkx",
+    "renfe",
+    "timetable"
+]
 
 # Structures
 
@@ -209,7 +274,7 @@ def station(name):
         return station
 
 
-def search(departure, arrival, departure_date, return_date=None):
+def search(departure, arrival, departure_date, return_date=None, via=None, age=27, systems=CT_SYSTEMS, cards=[]):
     "Search for a trip"
 
     headers = {
@@ -226,17 +291,17 @@ def search(departure, arrival, departure_date, return_date=None):
                                              {
                                                  "id": "d7c95386-9ed2-4366-8292-610d821940a3",
                                                  "label": "adult",
-                                                 "age": 27,
+                                                 "age": age,
                                                  "cards": [
-                                                     { "reference": "SNCF.Carte1225" }
+                                                     { "reference": card } for card in cards
                                                  ],
                                                  "cui": None
                                              }
                                          ],
-                                         "systems": ["sncf","db","busbud","idtgv","ouigo","trenitalia","ntv","hkx","renfe","timetable"],
+                                         "systems": systems if isinstance(systems, list) else [systems],
                                          "exchangeable_part": None,
                                          "departure_station_id": departure.id,
-                                         "via_station_id": None,
+                                         "via_station_id": via.id if via is not None else None,
                                          "arrival_station_id": arrival.id,
                                          "exchangeable_pnr_id": None
                                      }
@@ -248,11 +313,16 @@ def search(departure, arrival, departure_date, return_date=None):
 
 # Convinience functions
 
-def cheapest_trips(trips):
+def cheapest_trips(trips, max_price=None):
+    """Return cheapest trips availables
+    """
+
     min_trips = []
 
     for trip in trips:
-        if len(min_trips) == 0 or min_trips[0].price > trip.price:
+        if max_price is not None and trip.price > max_price:
+            continue
+        elif 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)

cheapest_ticket.py

@@ -5,46 +5,78 @@ import sys
 
 import captaintrain as ct
 
-def usage():
-    print("""cheapest_ticket.py departure-station arrival-station departure-date
-
-    departure-date: any string parsable by date -d
-""")
 
 if __name__ == "__main__":
-    if len(sys.argv) <= 3:
+    # Parse command line arguments
-        usage()
+    import argparse
-        sys.exit(1)
+    parser = argparse.ArgumentParser()
 
-    departure = ct.station(sys.argv[1])
+    parser.add_argument("-S", "--system", choices=ct.CT_SYSTEMS, action="append",
+                        help="backends in which search trips")
+
+    parser.add_argument("-A", "--age", type=int, default=27,
+                        help="your age")
+    parser.add_argument("-C", "--card", choices=ct.CT_CARDS, action="append", default=[],
+                        help="indicate cards you own")
+
+    parser.add_argument("departure_station", metavar='departure-station',
+                        help="name of the departure station")
+    parser.add_argument("via_station", metavar='via-station', nargs="?",
+                        help="name of a station you want on your trip")
+    parser.add_argument('arrival_station', metavar="arrival-station",
+                        help="name of the arrival station")
+    parser.add_argument('departure_date', metavar="departure-date", default="now",
+                        help="any string parsable by date -d")
+
+    args = parser.parse_args()
+
+    # Validate Systems
+    if args.system is None:
+        args.system = ct.CT_SYSTEMS
+
+    res = None
+    with open("rescaen.json") as f:
+        import json
+        res = ct.Search(**json.load(f))
+
+    # Search departure station
+    departure = ct.station(args.departure_station)
     if departure is None:
-        print("Unknown station '%s'" % sys.argv[1])
+        print("Unknown station '%s'" % args.departure_station)
         sys.exit(1)
+    print("From: %s (%s)" % (departure.name, departure.id))
 
-    arrival = ct.station(sys.argv[2])
+    # Search arrival station
+    arrival = ct.station(args.arrival_station)
     if arrival is None:
-        print("Unknown station '%s'" % sys.argv[2])
+        print("Unknown station '%s'" % args.arrival_station)
         sys.exit(1)
+    print("To: %s (%s)" % (arrival.name, arrival.id))
 
+    # Search via station
+    if args.via_station is not None:
+        via = ct.station(args.via_station)
+        if via is None:
+            print("Unknown station '%s'" % args.via_station)
+            sys.exit(1)
+        print("Via: %s (%s)" % (via.name, via.id))
+    else:
+        via = None
+
+    # Determine departure time
     departure_time = 0
-
+    with subprocess.Popen(["date", "-d", args.departure_date, "-u", "-Iseconds"], stdout=subprocess.PIPE) as f:
-    with subprocess.Popen(["date", "-d", sys.argv[3], "-u", "-Iseconds"], stdout=subprocess.PIPE) as f:
         departure_time = f.stdout.read().strip().decode()
 
     if departure_time == 0 or departure_time == "":
         sys.exit(1)
     else:
         departure_time = ct.parse_datetime(departure_time)
+    print("Departure on: %s" % departure_time)
 
-    print("From:", departure.name, departure.id)
-    print("To:", arrival.name, arrival.id)
-    print("Departure:", departure_time)
 
-    #with open("res.json") as f:
+    # Do things
-    #    import json
+    res = ct.search(departure, arrival, departure_time, via=via, systems=args.system, age=args.age, cards=args.card)
-    #    res = ct.Search(**json.load(f))
 
-    res = ct.search(departure, arrival, departure_time)
+    for trip in ct.cheapest_trips(res.trips):
-    min_trips = ct.cheapest_trips(res.trips)
-    for trip in min_trips:
         print(trip)