added progress bar for achiver

Serene-Arc · Jan 10, 2023 · 45840d7 · 45840d7
1 parent 059e661
commit 45840d7
Showing 1 changed file with 28 additions and 23 deletions.
diff --git a/bdfr/archiver.py b/bdfr/archiver.py
@@ -18,6 +18,7 @@
 from bdfr.configuration import Configuration
 from bdfr.connector import RedditConnector
 from bdfr.exceptions import ArchiverError
+from bdfr.progress_bar import Progress
 from bdfr.resource import Resource
 
 logger = logging.getLogger(__name__)
@@ -28,29 +29,33 @@ def __init__(self, args: Configuration):
         super(Archiver, self).__init__(args)
 
     def download(self):
-        for generator in self.reddit_lists:
-            try:
-                for submission in generator:
-                    try:
-                        if (submission.author and submission.author.name in self.args.ignore_user) or (
-                            submission.author is None and "DELETED" in self.args.ignore_user
-                        ):
-                            logger.debug(
-                                f"Submission {submission.id} in {submission.subreddit.display_name} skipped due to"
-                                f" {submission.author.name if submission.author else 'DELETED'} being an ignored user"
-                            )
-                            continue
-                        if submission.id in self.excluded_submission_ids:
-                            logger.debug(f"Object {submission.id} in exclusion list, skipping")
-                            continue
-                        logger.debug(f"Attempting to archive submission {submission.id}")
-                        self.write_entry(submission)
-                    except prawcore.PrawcoreException as e:
-                        logger.error(f"Submission {submission.id} failed to be archived due to a PRAW exception: {e}")
-            except prawcore.PrawcoreException as e:
-                logger.error(f"The submission after {submission.id} failed to download due to a PRAW exception: {e}")
-                logger.debug("Waiting 60 seconds to continue")
-                sleep(60)
+        progress = Progress(self.args.progress_bar)
+        with progress.context():
+            for generator in progress.wrap_subreddit_iter(self.reddit_lists):
+                try:
+                    for submission in progress.wrap_post_iter(generator):
+                        try:
+                            if (submission.author and submission.author.name in self.args.ignore_user) or (
+                                submission.author is None and "DELETED" in self.args.ignore_user
+                            ):
+                                logger.debug(
+                                    f"Submission {submission.id} in {submission.subreddit.display_name} skipped due to"
+                                    f" {submission.author.name if submission.author else 'DELETED'} being an ignored user"
+                                )
+                                continue
+                            if submission.id in self.excluded_submission_ids:
+                                logger.debug(f"Object {submission.id} in exclusion list, skipping")
+                                continue
+                            logger.debug(f"Attempting to archive submission {submission.id}")
+                            self.write_entry(submission)
+                            progress.post_log(submission, True)
+                        except prawcore.PrawcoreException as e:
+                            logger.error(f"Submission {submission.id} failed to be archived due to a PRAW exception: {e}")
+                            progress.post_log(submission, False)
+                except prawcore.PrawcoreException as e:
+                    logger.error(f"The submission after {submission.id} failed to download due to a PRAW exception: {e}")
+                    logger.debug("Waiting 60 seconds to continue")
+                    sleep(60)
 
     def get_submissions_from_link(self) -> list[list[praw.models.Submission]]:
         supplied_submissions = []