refactor: refactor screenshot_downloader, split to a clean functions

video_creation/screenshot_downloader.py

@@ -20,16 +20,23 @@ def download_screenshots_of_reddit_posts(reddit_object, screenshot_num, theme):
     with sync_playwright() as p:
         print_substep("Launching Headless Browser...")
 
-        browser = p.chromium.launch()
-        context = browser.new_context()
+        browser, context = p.chromium.launch(), browser.new_context()
+        page = context.new_page()
         
         if theme.casefold() == "dark":
             cookie_file = open('video_creation/cookies.json')
             cookies = json.load(cookie_file)
             context.add_cookies(cookies)
 
-        # Get the thread screenshot
-        page = context.new_page()
+        
+        save_thread_screenshot(page, reddit_object)
+        save_comments_screenshots(page, reddit_object, screenshot_num)
+
+        print_substep("Screenshots downloaded Successfully.",
+                      style="bold green")
+
+
+def save_thread_screenshot(page, reddit_object):
     page.goto(reddit_object["thread_url"])
     page.set_viewport_size(ViewportSize(width=1920, height=1080))
     if page.locator('[data-testid="content-gate"]').is_visible():
@@ -42,8 +49,11 @@ def download_screenshots_of_reddit_posts(reddit_object, screenshot_num, theme):
         path="assets/png/title.png"
     )
     
+    return page
+
+def save_comments_screenshots(page, reddit_object, screenshot_num):
     for idx, comment in track(
-            enumerate(reddit_object["comments"]), "Downloading screenshots..."
+            enumerate(reddit_object["comments"]), description="Downloading screenshots..."
         ):
 
             # Stop if we have reached the screenshot_num
@@ -57,6 +67,3 @@ def download_screenshots_of_reddit_posts(reddit_object, screenshot_num, theme):
             page.locator(f"#t1_{comment['comment_id']}").screenshot(
                 path=f"assets/png/comment_{idx}.png"
             )
-
-        print_substep("Screenshots downloaded Successfully.",
-                      style="bold green")