diff --git a/hyperdia.py b/hyperdia.py index fa4835f..d4f2f4e 100644 --- a/hyperdia.py +++ b/hyperdia.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 +import argparse from dataclasses import dataclass from datetime import datetime from itertools import zip_longest @@ -52,6 +53,17 @@ HEADERS = { } +def required_length(nmin, nmax): + class RequiredLength(argparse.Action): + def __call__(self, parser, args, values, option_string=None): + if not nmin <= len(values) <= nmax: + msg='argument "{f}" requires between {nmin} and {nmax} arguments'.format( + f=self.dest,nmin=nmin,nmax=nmax) + raise argparse.ArgumentTypeError(msg) + setattr(args, self.dest, values) + return RequiredLength + + @dataclass class HyperdiaStep: @@ -298,6 +310,7 @@ def hyperdia_search(start_station: str, end_station: str, hour: int, year: int = 2020, max_route: int = 5, via: List[str] = None) -> List[str]: + # TODO: Error checking raw_result = get_hyperdia_data(start_station, end_station, hour, minute, day, month, year, max_route, via) @@ -306,3 +319,36 @@ def hyperdia_search(start_station: str, end_station: str, hour: int, for trip in results: print(trip_summary(trip)) + + +def main(): + + parser = argparse.argumentParser() + parser.add_argument("-t", "--time", help="Hour of travel", + type=lambda d: datetime.strptime(d, '%H.%M').time()) + parser.add_argument("-d", "--date", help="Date of travel", + type=lambda d: datetime.strptime(d, "%Y-%m-%d").date()) + parser.add_arguments("--max-routes", help="Maximum number of routes", + type=int) + parser.add_argument("--via", nargs='+', action=required_length(1, 3), + help="Stations to force route through (min 1, max 3)") + parser.add_argument("start_station", help="Start station") + parser.add_argument("end_station", help="End station") + + options = parser.parse_args() + + hour, minute = options.time.hour, options.time.minute + day, month, year = options.date.day, options.date.month, options.date.year + + if month > 9: + # Add "0" in front of single-digit months + month = str(month).zfill(2) + else: + month = str(month) + + hyperdia_search(options.start_station, options.end_station, hour, minute, + day, month, year, options.max_routes, via=options.via) + + +if __name__ == "__main__": + main()