Use argparse and improve search

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:
-        usage()
-        sys.exit(1)
+    # Parse command line arguments
+    import argparse
+    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", sys.argv[3], "-u", "-Iseconds"], stdout=subprocess.PIPE) as f:
+    with subprocess.Popen(["date", "-d", args.departure_date, "-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:
-    #    import json
-    #    res = ct.Search(**json.load(f))
+    # Do things
+    res = ct.search(departure, arrival, departure_time, via=via, systems=args.system, age=args.age, cards=args.card)
 
-    res = ct.search(departure, arrival, departure_time)
-    min_trips = ct.cheapest_trips(res.trips)
-    for trip in min_trips:
+    for trip in ct.cheapest_trips(res.trips):
         print(trip)