Add get_mega_millions_object function

util/drawing.py

@@ -5,7 +5,7 @@
 """
 from datetime import datetime, timedelta
 
-from data.database import PowerballDraw
+from data.database import PowerballDraw, MegaMillionsDraw
 
 
 def get_next_draw_date(game: str | None, from_date: str | None = None) -> str | None:
@@ -55,7 +55,7 @@ def get_powerball_object(draw_date: str, draw_result: dict) -> PowerballDraw:
     
     Examples:
     >>> get_powerball_object(draw_date='2026-02-28', draw_result={'main_balls': [6, 20, 35, 54, 65], 'powerball': 10, 'multiplier': 4})
-    PowerballDraw(draw_date='2026-02-28', main_ball1=6, main_ball2=20, main_ball3=35, main_ball4=54, main_ball5=65, powerball=10, power_play=4)
+    PowerballDraw(draw_date=datetime(2026, 2, 28), main_ball1=6, main_ball2=20, main_ball3=35, main_ball4=54, main_ball5=65, powerball=10, power_play=4)
     """
     draw_date = datetime.strptime(draw_date, '%Y-%m-%d').date()
     mb1, mb2, mb3, mb4, mb5 = draw_result['main_balls']
@@ -63,4 +63,24 @@ def get_powerball_object(draw_date: str, draw_result: dict) -> PowerballDraw:
     power_play = draw_result['multiplier']
 
     return PowerballDraw(draw_date=draw_date,main_ball1=mb1, main_ball2=mb2, main_ball3=mb3, main_ball4=mb4,
-         main_ball5=mb5, powerball=pb, power_play=power_play)
+        main_ball5=mb5, powerball=pb, power_play=power_play)
+
+
+def get_mega_millions_object(draw_date: str, draw_result: dict) -> MegaMillionsDraw:
+    """
+    :param draw_date: date for the drawing data
+    :param draw_result: drawing results for the draw_date
+    :return: MegaMillionsDraw object
+
+    Examples:
+    >>> get_mega_millions_object(draw_date='2010-02-02', draw_result={'main_balls': [6, 7, 26, 27, 49], 'mega-ball': 9, 'megaplier': 4})
+    MegaMillionsDraw(draw_date=date, main_ball1=6, main_ball2=7, main_ball3=26, main_ball4=27, main_ball5=49, mega_ball=9, megaplier=4)
+    """
+    print(draw_result)
+    draw_date = datetime.strptime(draw_date, '%Y-%m-%d').date()
+    mb1, mb2, mb3, mb4, mb5 = draw_result['main_balls']
+    mega = draw_result['mega-ball']
+    megaplier = draw_result['multiplier']
+
+    return MegaMillionsDraw(draw_date=draw_date, main_ball1=mb1, main_ball2=mb2, main_ball3=mb3, main_ball4=mb4,
+        main_ball5=mb5, mega_ball=mega, megaplier=megaplier)