Simplify get_next_draw_date

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
+    return from_date.strftime(return_date_format)
-    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