diff --git a/util/drawing.py b/util/drawing.py index 32993ae..88d1f29 100644 --- a/util/drawing.py +++ b/util/drawing.py @@ -1,8 +1,8 @@ -from datetime import datetime +from datetime import datetime, timedelta from dateutil.relativedelta import relativedelta, MO, TU, WE, FR, SA -def get_next_draw_date(game: str | None, from_date: str | None = None) -> str: +def get_next_draw_date(game: str | None, from_date: str | None = None) -> str | None: """ Gets the next drawing date from the current or specified date for the given game If the current date is the next drawing date, the current date is returned @@ -17,33 +17,25 @@ def get_next_draw_date(game: str | None, from_date: str | None = None) -> str: >>> get_next_draw_date(game='powerball', from_date='02/15/2026') '02/16/2026' """ + target_days = { + "powerball": [0, 2, 5], + "megamillions": [1, 4] + } + return_date = None return_date_format = '%m/%d/%Y' + + # if a valid game is not specified, return None + if game is None or game not in target_days.keys(): + return return_date + # if no date was specified, use today's date if from_date is None: from_date = datetime.today() else: from_date = datetime.strptime(from_date, return_date_format) - day_of_week = from_date.weekday() + while from_date.weekday() not in target_days[game]: + from_date += timedelta(days=1) - # target the selected game - if game == 'megamillions': - """drawings are on Tuesday and Friday""" - if day_of_week in [1, 4]: - return_date = from_date.strftime(return_date_format) - elif day_of_week in [2, 3]: - return_date = (from_date + relativedelta(weekday=FR)).strftime(return_date_format) - elif day_of_week in [5, 6, 0]: - return_date = (from_date + relativedelta(weekday=TU)).strftime(return_date_format) - elif game == 'powerball': - """drawings are on Monday, Wednesday, and Saturday""" - if day_of_week in [0, 2, 5]: - return_date = from_date.strftime(return_date_format) - elif day_of_week == 1: - return_date = (from_date + relativedelta(weekday=WE)).strftime(return_date_format) - elif day_of_week in [3, 4]: - return_date = (from_date + relativedelta(weekday=SA)).strftime(return_date_format) - elif day_of_week == 6: - return_date = (from_date + relativedelta(weekday=MO)).strftime(return_date_format) - return return_date + return from_date.strftime(return_date_format)