From 14b7f574b778e7b586e8755dc1e071ece4dc0603 Mon Sep 17 00:00:00 2001
From: Tomovic <jovantomovic5@gmail.com>
Date: Thu, 7 Jul 2022 18:10:16 +0200
Subject: [PATCH] add: NSFW to videos.json

---
 utils/videos.py                         | 4 +++-
 video_creation/final_video.py           | 6 +++++-
 video_creation/screenshot_downloader.py | 2 ++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/utils/videos.py b/utils/videos.py
index d050ab8..9f627ec 100755
--- a/utils/videos.py
+++ b/utils/videos.py
@@ -34,7 +34,7 @@ def check_done(
     return redditobj
 
 
-def save_data(filename: str, reddit_title: str, reddit_id: str, credit: str):
+def save_data(filename: str, reddit_title: str, reddit_id: str, credit: str, NSFW: bool):
     """Saves the videos that have already been generated to a JSON file in video_creation/data/videos.json
 
     Args:
@@ -42,6 +42,7 @@ def save_data(filename: str, reddit_title: str, reddit_id: str, credit: str):
         @param filename:
         @param reddit_id:
         @param reddit_title:
+        @param NSFW:
     """
     with open("./video_creation/data/videos.json", "r+", encoding="utf-8") as raw_vids:
         done_vids = json.load(raw_vids)
@@ -53,6 +54,7 @@ def save_data(filename: str, reddit_title: str, reddit_id: str, credit: str):
             "background_credit": credit,
             "reddit_title": reddit_title,
             "filename": filename,
+            "nsfw": NSFW,
         }
         done_vids.append(payload)
         raw_vids.seek(0)
diff --git a/video_creation/final_video.py b/video_creation/final_video.py
index d08e4cb..0f01b6c 100755
--- a/video_creation/final_video.py
+++ b/video_creation/final_video.py
@@ -28,6 +28,10 @@ console = Console()
 
 W, H = 1080, 1920
 
+nsfw = False
+def isNSFW():
+    global nsfw
+    nsfw = True
 
 def name_normalize(name: str) -> str:
     name = re.sub(r'[?\\"%*:|<>]', "", name)
@@ -115,7 +119,7 @@ def make_final_video(
     filename = f"{name_normalize(title)}.mp4"
     subreddit = settings.config["reddit"]["thread"]["subreddit"]
 
-    save_data(filename, title, idx, background_config[2])
+    save_data(filename, title, idx, background_config[2], nsfw)
 
     if not exists(f"./results/{subreddit}"):
         print_substep("The results folder didn't exist so I made it")
diff --git a/video_creation/screenshot_downloader.py b/video_creation/screenshot_downloader.py
index 6fb9ef4..fc83d79 100644
--- a/video_creation/screenshot_downloader.py
+++ b/video_creation/screenshot_downloader.py
@@ -12,6 +12,7 @@ from rich.progress import track
 import translators as ts
 
 from utils.console import print_step, print_substep
+from video_creation.final_video import isNSFW
 
 storymode = False
 
@@ -47,6 +48,7 @@ def download_screenshots_of_reddit_posts(reddit_object: dict, screenshot_num: in
         if page.locator('[data-testid="content-gate"]').is_visible():
             # This means the post is NSFW and requires to click the proceed button.
 
+            isNSFW()
             print_substep("Post is NSFW. You are spicy...")
             page.locator('[data-testid="content-gate"] button').click()
             page.locator(