You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
RedditVideoMakerBot/video_creation/screenshot_downloader.py

73 lines
2.7 KiB

#!/usr/bin/env python3
from playwright.sync_api import sync_playwright, ViewportSize
from pathlib import Path
from rich.progress import track
from utils.console import print_step, print_substep
import json
from rich.console import Console
console = Console()
def download_screenshots_of_reddit_posts(reddit_object, screenshot_num, theme):
"""Downloads screenshots of reddit posts as they are seen on the web.
Args:
reddit_object: The Reddit Object you received in askreddit.py
screenshot_num: The number of screenshots you want to download.
"""
print_step("Downloading screenshots of reddit posts...")
# ! Make sure the reddit screenshots folder exists
Path("assets/png").mkdir(parents=True, exist_ok=True)
with sync_playwright() as p:
print_substep("Launching Headless Browser...")
browser = p.chromium.launch()
context = browser.new_context()
try:
if theme.casefold() == "dark":
cookie_file = open('video_creation/cookies.json')
cookies = json.load(cookie_file)
context.add_cookies(cookies)
except AttributeError:
pass
# 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()
page.locator('[data-click-id="text"] button').click() # Remove "Click to see nsfw" Button in Screenshot
page.locator('[data-test-id="post-content"]').screenshot(
path="assets/png/title.png"
)
for idx, comment in track(
enumerate(reddit_object["comments"])
):
#allow user to see what comment is being saved
print_substep(f"Downloading screenshot {idx + 1}")
# 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()
page.goto(f'https://reddit.com{comment["comment_url"]}')
page.locator(f"#t1_{comment['comment_id']}").screenshot(
path=f"assets/png/comment_{idx}.png"
)
#let user know that the screenshots are done
console.log(f"[bold green]Saved {idx + 1} screenshots.")