diff --git a/video_creation/screenshot_downloader.py b/video_creation/screenshot_downloader.py index 18fe5c5..3aca45f 100644 --- a/video_creation/screenshot_downloader.py +++ b/video_creation/screenshot_downloader.py @@ -3,6 +3,7 @@ from pathlib import Path from playwright.sync_api import sync_playwright, ViewportSize from rich.progress import track +from rich.console import Console from utils.console import print_step, print_substep @@ -14,46 +15,48 @@ def download_screenshots_of_reddit_posts(reddit_object, screenshot_num, theme): reddit_object: The Reddit Object you received in askreddit.py screenshot_num: The number of screenshots you want to download. """ + console = Console() + print_step("Downloading Screenshots of Reddit Posts 📷") # ! Make sure the reddit screenshots folder exists Path("assets/png").mkdir(parents=True, exist_ok=True) - print_substep("Launching Headless Browser...") - with sync_playwright() as browser_: - browser = browser_.chromium.launch() - context = browser.new_context() + with console.status("[bold]Launching Headless Browser ...", spinner="simpleDots"): + with sync_playwright() as browser_: + browser = browser_.chromium.launch() + context = browser.new_context() - if theme.casefold() == "dark": - with open("video_creation/cookies.json", encoding="utf-8") as cookie_file: - cookies = json.load(cookie_file) - context.add_cookies(cookies) + if theme.casefold() == "dark": + with open("video_creation/cookies.json", encoding="utf-8") as cookie_file: + cookies = json.load(cookie_file) + context.add_cookies(cookies) - # Get the thread screenshot - page = context.new_page() - page.goto(reddit_object["thread_url"]) - page.set_viewport_size(ViewportSize(width=1920, height=1080)) - if page.locator('[data-testid="content-gate"]').is_visible(): - # This means the post is NSFW and requires to click the proceed button. + # Get the thread screenshot + page = context.new_page() + page.goto(reddit_object["thread_url"]) + page.set_viewport_size(ViewportSize(width=1920, height=1080)) + if page.locator('[data-testid="content-gate"]').is_visible(): + # This means the post is NSFW and requires to click the proceed button. - print_substep("Post is NSFW. You are spicy...") - page.locator('[data-testid="content-gate"] button').click() + print_substep("Post is NSFW. You are spicy...") + page.locator('[data-testid="content-gate"] button').click() - page.locator('[data-test-id="post-content"]').screenshot(path="assets/png/title.png") + page.locator('[data-test-id="post-content"]').screenshot(path="assets/png/title.png") - for idx, comment in track( - enumerate(reddit_object["comments"]), "Downloading screenshots..." - ): - # Stop if we have reached the screenshot_num - if idx >= screenshot_num: - break + for idx, comment in track( + enumerate(reddit_object["comments"]), "Downloading screenshots..." + ): + # Stop if we have reached the screenshot_num + if idx >= screenshot_num: + break - if page.locator('[data-testid="content-gate"]').is_visible(): - page.locator('[data-testid="content-gate"] button').click() + if page.locator('[data-testid="content-gate"]').is_visible(): + page.locator('[data-testid="content-gate"] button').click() - page.goto(f'https://reddit.com{comment["comment_url"]}') - page.locator( - f"#t1_{comment['comment_id']}" - ).screenshot(path=f"assets/png/comment_{idx}.png") + page.goto(f'https://reddit.com{comment["comment_url"]}') + page.locator( + f"#t1_{comment['comment_id']}" + ).screenshot(path=f"assets/png/comment_{idx}.png") print_substep("Screenshots downloaded Successfully.", style="bold green")