Rewrite main function and add support functions
Add argparse argument handling
This commit is contained in:
chris committed
1 parent
39d5853638
commit
31a98bfb64
1 file changed
+209
-114
@@ -1,130 +1,225 @@
|
||||
from datetime import datetime
|
||||
import argparse
|
||||
from datetime import date, datetime
|
||||
import random
|
||||
from time import sleep
|
||||
|
||||
from sqlalchemy import func
|
||||
from sqlalchemy.exc import IntegrityError, PendingRollbackError
|
||||
from tqdm import tqdm
|
||||
|
||||
from data.database import dal, MegaMillionsDraw, PowerballDraw
|
||||
import lottery_predictor.collect as predictor
|
||||
from util.drawing import get_mega_millions_object, get_powerball_object
|
||||
from util.environment import load_environment_variables
|
||||
from data.database import dal2
|
||||
import lottery_predictor.analyze as analyze
|
||||
from lottery_predictor.config import PROJECT_VARIABLES, GAME_INFO
|
||||
from util.scrape import scrape_game_history
|
||||
|
||||
project_variables = load_environment_variables()
|
||||
BAR_FORMAT = PROJECT_VARIABLES["LP_BAR_FORMAT"]
|
||||
|
||||
def dry_run_fill_game_data(game: str, year: int | None = None) -> None:
|
||||
"""
|
||||
|
||||
:param game: Name of the game to process
|
||||
:param year: A starting year or None
|
||||
:return: Nothing
|
||||
|
||||
A dummy job to go through the motions of updating without doing anything
|
||||
"""
|
||||
if year is None:
|
||||
year = GAME_INFO[game]['start_year']
|
||||
elif not isinstance(year, int):
|
||||
raise ValueError(f"Year must be a valid year or None, received {year}")
|
||||
elif year not in range(1992, 2030):
|
||||
raise ValueError(f"Year must be within the range 1992 and 2030, received {year}")
|
||||
|
||||
for year in range(year, datetime.today().year + 1):
|
||||
for _ in tqdm(
|
||||
range(random.randint(1, 106)),
|
||||
desc=f"\tProcessing year {year}",
|
||||
bar_format=BAR_FORMAT
|
||||
):
|
||||
sleep(0.25)
|
||||
|
||||
if year != datetime.now().year:
|
||||
for _ in tqdm(
|
||||
range(5),
|
||||
desc='\tPseudo API backoff',
|
||||
bar_format=BAR_FORMAT
|
||||
):
|
||||
sleep(1)
|
||||
else:
|
||||
print()
|
||||
|
||||
|
||||
def fill_game_data(game: str, year: int | None = None) -> None:
|
||||
"""
|
||||
:param game: Name of the game to process
|
||||
:param year: A starting year or None
|
||||
:return: Nothing
|
||||
|
||||
If the year is None, all game data will be scraped for each year the game has data available.
|
||||
Otherwise, the scrape will be limited to a range of years starting with the year specified
|
||||
Not all data scraped will be used for predictions, but could be useful later.
|
||||
"""
|
||||
if year is None:
|
||||
year = GAME_INFO[game]['start_year']
|
||||
elif not isinstance(year, int):
|
||||
raise ValueError(f"Year must be a valid year or None, received {year}")
|
||||
elif year not in range(1992, 2030):
|
||||
raise ValueError(f"Year must be within the range 1992 and 2030, received {year}")
|
||||
|
||||
# scrape records by year and insert into the database
|
||||
for year in range(year, datetime.today().year + 1):
|
||||
scrape_data = scrape_game_history(game=GAME_INFO[game]['scrape_name'], year=year)
|
||||
# convert and insert records
|
||||
for _ in tqdm(
|
||||
range(len(scrape_data)),
|
||||
desc=f"Processing year {year}",
|
||||
bar_format=BAR_FORMAT
|
||||
):
|
||||
for draw_date, numbers in scrape_data.items():
|
||||
record = GAME_INFO[game]['object_function'](draw_date=draw_date, draw_result=numbers)
|
||||
dal2.add(record)
|
||||
# sleep for 30 seconds to avoid overtaxing the endpoint, if the processed year is not the current one
|
||||
if year != datetime.now().year:
|
||||
for _ in tqdm(
|
||||
range(30),
|
||||
desc='\tAPI backoff',
|
||||
bar_format=BAR_FORMAT
|
||||
):
|
||||
sleep(1)
|
||||
|
||||
|
||||
def update_games(dry_run: bool = False) -> None:
|
||||
"""
|
||||
:return: None
|
||||
|
||||
Loop over games in the game_info dictionary and update the stored records
|
||||
"""
|
||||
date_format = PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"]
|
||||
|
||||
for game, info in GAME_INFO.items():
|
||||
print(f"Updating {game} data...\n")
|
||||
if dal2.count(table_name=info['table']) == 0:
|
||||
if not dry_run:
|
||||
fill_game_data(game=game)
|
||||
else:
|
||||
dry_run_fill_game_data(game=game)
|
||||
else:
|
||||
most_recent = dal2.most_recent(table_name=info['table'])
|
||||
start_year = datetime.strptime(str(most_recent[0]), date_format).year
|
||||
if not dry_run:
|
||||
fill_game_data(game=game, year=start_year)
|
||||
else:
|
||||
dry_run_fill_game_data(game=game, year=start_year)
|
||||
|
||||
|
||||
def get_table_counts(from_date: date | str | None) -> dict[str, int]:
|
||||
"""
|
||||
Loops over the tables in the database and gets a record count for each one.
|
||||
|
||||
:return: A dictionary of table names and record counts
|
||||
"""
|
||||
counts = {}
|
||||
for game, info in GAME_INFO.items():
|
||||
if isinstance(from_date, date):
|
||||
counts[game] = dal2.count(table_name=info['table'], from_date=from_date)
|
||||
elif from_date == 'rule-change':
|
||||
rule_date = datetime.strptime(info['rule_change'], PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"])
|
||||
counts[game] = dal2.count(table_name=info['table'], from_date=rule_date)
|
||||
else:
|
||||
counts[game] = dal2.count(table_name=info['table'])
|
||||
return counts
|
||||
|
||||
|
||||
def get_prediction(game: str, window_size: int, test_case: bool = False) -> str:
|
||||
# Get the list of drawings to make the prediction from
|
||||
from_date = GAME_INFO[game]['rule_change']
|
||||
|
||||
if test_case:
|
||||
if window_size == -1:
|
||||
drawings = analyze.load_dataframe_by_dates(game=game, start_date=from_date)
|
||||
else:
|
||||
drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size)
|
||||
window_size = len(drawings) if window_size == -1 else window_size
|
||||
test_case = drawings.iloc[0]
|
||||
game_ball_out = 'powerball' if game == 'Powerball' else 'mega_ball'
|
||||
case_string = (f"Draw Date: {test_case['draw_date']}, "
|
||||
f"Main Balls: [{test_case['main_ball1']}, {test_case['main_ball2']}, {test_case['main_ball3']}, "
|
||||
f"{test_case['main_ball4']}, {test_case['main_ball5']}], Game Ball: {test_case[game_ball_out]}")
|
||||
|
||||
predicted = analyze.make_prediction(data_frame=drawings.iloc[1:], window_size=(window_size - 2))
|
||||
main_balls = ", ".join(predicted[0].astype(str))
|
||||
game_ball = "".join(predicted[1].astype(str))
|
||||
return f"{case_string}\n Prediction: Main Balls: [{main_balls}], Game Ball: {game_ball}\n"
|
||||
else:
|
||||
if window_size == -1:
|
||||
drawings = analyze.load_dataframe_by_dates(game=game, start_date=from_date)
|
||||
else:
|
||||
drawings = analyze.load_dataframe_most_recent(game=game, limit=window_size)
|
||||
window_size = len(drawings) if window_size == -1 else window_size
|
||||
predicted = analyze.make_prediction(data_frame=drawings, window_size=(window_size - 1))
|
||||
main_balls = ", ".join(predicted[0].astype(str))
|
||||
game_ball = "".join(predicted[1].astype(str))
|
||||
return f"Prediction: Main Balls: {main_balls}, Game Ball: {game_ball}\n"
|
||||
|
||||
|
||||
def main():
|
||||
print("Lottery Analyzer initializing...")
|
||||
# convert specific variables for use
|
||||
mm_start_year = int(project_variables["LP_MEGA_MILLIONS_INCEPTION"])
|
||||
pb_start_year = int(project_variables["LP_POWERBALL_INCEPTION"])
|
||||
end_year = datetime.today().year
|
||||
date_format = project_variables["LP_DATE_INSERT_FORMAT"]
|
||||
max_mm_main = int(project_variables["LP_MEGA_MILLIONS_MAX_MAIN_BALL"])
|
||||
max_megaball = int(project_variables["LP_MEGA_MILLIONS_MAX_MEGA_BALL"])
|
||||
max_pb_main = int(project_variables["LP_POWERBALL_MAX_MAIN_BALL"])
|
||||
max_powerball = int(project_variables["LP_POWERBALL_MAX_POWERBALL"])
|
||||
parser = argparse.ArgumentParser(description="A lottery prediction tool.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
update_parser = subparsers.add_parser("update", help="Update database records")
|
||||
update_parser.add_argument("-y", "--year", type=int, help="Year to start updates with")
|
||||
update_parser.add_argument(
|
||||
"-d", "--dryrun",
|
||||
action="store_true",
|
||||
help="Run the update process without actually making any database or API calls"
|
||||
)
|
||||
|
||||
# connect to the database
|
||||
print("Connecting to the database...")
|
||||
dal.connect()
|
||||
session = dal.Session()
|
||||
record_parser = subparsers.add_parser("record-count", help="Get record counts by table in the database")
|
||||
record_parser.add_argument(
|
||||
"-f", "--from-date",
|
||||
type=str,
|
||||
help="Date from which to base the count on (ex. 2026-01-01). \
|
||||
The default value is rule-change for the date when game rules last changed"
|
||||
)
|
||||
|
||||
print("Collecting and processing missing Mega Millions draws...")
|
||||
# if there are no records in the MegaMillionsDraw table...
|
||||
if predictor.get_record_count('MegaMillionsDraw') == 0:
|
||||
# loop over all years and scrape the data for each
|
||||
for year in range(mm_start_year, end_year + 1):
|
||||
print(f"\tCollecting {year}...")
|
||||
scrape_data = scrape_game_history(game='mega-millions', year=year)
|
||||
# convert and insert the records
|
||||
for key, value in scrape_data.items():
|
||||
record = get_mega_millions_object(draw_date=key, draw_result=value)
|
||||
try:
|
||||
session.add(record)
|
||||
session.commit()
|
||||
except IntegrityError:
|
||||
pass
|
||||
except PendingRollbackError:
|
||||
session.rollback()
|
||||
finally:
|
||||
session.close()
|
||||
# sleep between records and years to avoid querying the site too often
|
||||
sleep(60)
|
||||
predict_parser = subparsers.add_parser("predict", help="Predict the next drawing")
|
||||
predict_parser.add_argument(
|
||||
"--game",
|
||||
type=str,
|
||||
required=True,
|
||||
help="The game to predict the next result for [MegaMillions, Powerball]"
|
||||
)
|
||||
predict_parser.add_argument(
|
||||
"--window-size",
|
||||
type=int,
|
||||
default=-1,
|
||||
help="The number of records to use in the prediction model, defaults to -1 for all records"
|
||||
)
|
||||
predict_parser.add_argument(
|
||||
"--test",
|
||||
action="store_true",
|
||||
help="Use the most recent drawing as a test subject to find the right window_size"
|
||||
)
|
||||
|
||||
# commit all the records to the database
|
||||
args = parser.parse_args()
|
||||
match args.command:
|
||||
case "record-count":
|
||||
# load the date format for the database
|
||||
date_format = PROJECT_VARIABLES["LP_DATE_INSERT_FORMAT"]
|
||||
if args.from_date:
|
||||
# if args.from_date is 'rule-change' use that date from vars, otherwise use the date provided
|
||||
if args.from_date != 'rule-change':
|
||||
from_date = datetime.strptime(args.from_date, date_format)
|
||||
# get the table row counts from the date specified
|
||||
table_counts = get_table_counts(from_date=from_date)
|
||||
else:
|
||||
most_recent = session.query(func.max(MegaMillionsDraw.draw_date)).first()
|
||||
start_year = datetime.strptime(str(most_recent[0]), date_format).year
|
||||
# loop over the years needed and gather records
|
||||
for year in range(start_year, end_year + 1):
|
||||
print(f"\tCollecting {year}...")
|
||||
scrape_data = scrape_game_history(game='powerball', year=year)
|
||||
# convert and insert the records
|
||||
for key, value in scrape_data.items():
|
||||
record = get_powerball_object(draw_date=key, draw_result=value)
|
||||
try:
|
||||
session.add(record)
|
||||
session.commit()
|
||||
except IntegrityError:
|
||||
pass
|
||||
except PendingRollbackError:
|
||||
session.rollback()
|
||||
finally:
|
||||
session.close()
|
||||
# sleep between records and years to avoid querying the site too often
|
||||
sleep(60)
|
||||
|
||||
print("Collecting and processing missing Powerball draws...")
|
||||
# if there are no records in the PowerBallDraw table...
|
||||
if len(session.query(PowerballDraw).all()) == 0:
|
||||
# loop over all years and scrape the data for each
|
||||
for year in range(pb_start_year, end_year + 1):
|
||||
print(f"\tCollecting {year}...")
|
||||
scrape_data = scrape_game_history(game='powerball', year=year)
|
||||
# convert and insert the records
|
||||
for key, value in scrape_data.items():
|
||||
record = get_powerball_object(draw_date=key, draw_result=value)
|
||||
try:
|
||||
session.add(record)
|
||||
session.commit()
|
||||
except IntegrityError:
|
||||
pass
|
||||
except PendingRollbackError:
|
||||
session.rollback()
|
||||
finally:
|
||||
session.close()
|
||||
# sleep between records and years to avoid querying the site too often
|
||||
sleep(60)
|
||||
|
||||
# commit all the records to the database
|
||||
session.commit()
|
||||
table_counts = get_table_counts(from_date=args.from_date)
|
||||
else:
|
||||
most_recent = session.query(func.max(PowerballDraw.draw_date)).first()
|
||||
start_year = datetime.strptime(str(most_recent[0]), str(date_format)).year
|
||||
# loop over the years needed and gather records
|
||||
for year in range(start_year, end_year + 1):
|
||||
print(f"\tCollecting {year}...")
|
||||
scrape_data = scrape_game_history(game='powerball', year=year)
|
||||
# convert and insert the records
|
||||
for key, value in scrape_data.items():
|
||||
record = get_powerball_object(draw_date=key, draw_result=value)
|
||||
try:
|
||||
session.add(record)
|
||||
session.commit()
|
||||
except IntegrityError:
|
||||
pass
|
||||
except PendingRollbackError:
|
||||
session.rollback()
|
||||
finally:
|
||||
session.close()
|
||||
# sleep between records and years to avoid querying the site too often
|
||||
sleep(60)
|
||||
|
||||
mm_count = session.query(MegaMillionsDraw).count()
|
||||
pb_count = session.query(PowerballDraw).count()
|
||||
|
||||
print(f"There are {mm_count} Mega Millions draw records and {pb_count} Powerball records")
|
||||
table_counts = get_table_counts(from_date=None)
|
||||
for table, count in table_counts.items():
|
||||
print(f"{table}:\t{count}")
|
||||
case "update":
|
||||
update_games(dry_run=args.dryrun)
|
||||
case "predict":
|
||||
print(get_prediction(game=args.game, window_size=args.window_size, test_case=args.test))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in new issue
Block a user