From e1ecf4fd7a7fdd95344d1651475b1e99c5e3eaba Mon Sep 17 00:00:00 2001
From: iaacornus <iaacornus.devel@gmail.com>
Date: Wed, 8 Jun 2022 00:10:24 +0800
Subject: [PATCH] allow users to specify how many of the comments should be
 included

---
 cli.py              |  7 +++++++
 main.py             | 14 ++++++++++++--
 reddit/subreddit.py |  7 ++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/cli.py b/cli.py
index 021eba1..0faba9e 100644
--- a/cli.py
+++ b/cli.py
@@ -47,6 +47,12 @@ def program_options():
         help="Use the given thread link instead of randomized.",
         action="store"
     )
+    parser.add_argument(
+        "-n",
+        "--number",
+        help="Number of comments to include.",
+        action="store"
+    )
     parser.add_argument(
         "--setup",
         "--setup",
@@ -65,6 +71,7 @@ def program_options():
                     args.background,
                     args.filename,
                     args.thread,
+                    args.number,
                 )
                 if not create:
                     try_again = input("Something went wrong! Try again? [y/N] > ").strip()
diff --git a/main.py b/main.py
index 7f55595..d470aa1 100644
--- a/main.py
+++ b/main.py
@@ -13,7 +13,13 @@ from utils.console import print_markdown, print_substep
 from setup_program import setup
 
 
-def main(subreddit_=None, background=None, filename=None, thread_link_=None):
+def main(
+        subreddit_=None,
+        background=None,
+        filename=None,
+        thread_link_=None,
+        number_of_comments=None
+    ):
     """
     Load .env file if exists. If it doesnt exist, print a warning and
     launch the setup wizard. If there is a .env file, check if the
@@ -53,7 +59,11 @@ def main(subreddit_=None, background=None, filename=None, thread_link_=None):
     if configured:
         print_substep("Enviroment Variables are set! Continuing ...", style_="bold green")
 
-        reddit_object = get_subreddit_threads(subreddit_, thread_link_)
+        reddit_object = get_subreddit_threads(
+            subreddit_,
+            thread_link_,
+            number_of_comments
+        )
         length, number_of_comments = save_text_to_mp3(reddit_object)
         download_screenshots_of_reddit_posts(
             reddit_object,
diff --git a/reddit/subreddit.py b/reddit/subreddit.py
index c3042ed..0eb953b 100644
--- a/reddit/subreddit.py
+++ b/reddit/subreddit.py
@@ -15,7 +15,7 @@ from rich.console import Console
 from utils.console import print_step, print_substep
 
 
-def get_subreddit_threads(subreddit_, thread_link_):
+def get_subreddit_threads(subreddit_, thread_link_, number_of_comments):
     """
     Takes subreddit_ as parameter which defaults to None, but in this
     case since it is None, it would raise ValueError, thus defaulting
@@ -109,7 +109,12 @@ def get_subreddit_threads(subreddit_, thread_link_):
         content["thread_post"] = submission.selftext
         content["comments"] = []
 
+        comment_count = 0
         for top_level_comment in submission.comments:
+            if number_of_comments is not None:
+                if comment_count > number_of_comments:
+                    break
+
             if not top_level_comment.stickied:
                 content["comments"].append(
                     {