Use argparse and improve search

This commit is contained in:
nemunaire 2016-07-10 22:25:41 +02:00
parent c9618a9256
commit 7c13541e3a
2 changed files with 132 additions and 30 deletions

View File

@ -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)

View File

@ -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)