diff --git a/main.py b/main.py index fafc485..1f66175 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ import argparse import random +import sys from datetime import date, datetime from time import sleep @@ -219,6 +220,10 @@ def get_prediction( return prediction_output +def check_game_name(game: str) -> bool: + return game in GAME_INFO + + def main(): parser = argparse.ArgumentParser(description="A lottery prediction tool.") subparsers = parser.add_subparsers(dest="command", required=True) @@ -258,8 +263,9 @@ def main(): ticket_parser.add_argument( "--game", type=str, + choices=[game for game in GAME_INFO.keys()], required=True, - help="The game to predict the next result for [MegaMillions, Powerball]", + help="The game to predict the next result for", ) predict_parser = subparsers.add_parser( @@ -269,8 +275,9 @@ def main(): predict_parser.add_argument( "--game", type=str, + choices=[game for game in GAME_INFO.keys()], required=True, - help="The game to predict the next result for [MegaMillions, Powerball]", + help="The game to predict the next result for", ) predict_parser.add_argument( "--window-size", @@ -289,6 +296,9 @@ def main(): args = parser.parse_args() match args.command: case "generate-ticket": + if not check_game_name(args.game): + print(f"{args.game} is not a valid game name.") + sys.exit() main_balls, game_ball = ticket.generate_random_ticket( max_main=GAME_INFO[args.game]['max_main_ball'], max_game=GAME_INFO[args.game]['max_game_ball'], @@ -317,6 +327,9 @@ def main(): case "update": update_games(dry_run=args.dryrun) case "predict": + if not check_game_name(args.game): + print(f"{args.game} is not a valid game name.") + sys.exit() if args.test: recent, prediction = get_prediction( game=args.game,