progress bar refactored to a class

Author: stared

bdfr/downloader.py

@@ -14,16 +14,14 @@
 import praw.exceptions
 import praw.models
 import prawcore
-from tqdm import tqdm
-from tqdm.contrib.logging import logging_redirect_tqdm
 
 from bdfr import exceptions as errors
 from bdfr.configuration import Configuration
 from bdfr.connector import RedditConnector
+from bdfr.progress_bar import Progress
 from bdfr.site_downloaders.download_factory import DownloadFactory
 
 logger = logging.getLogger(__name__)
-logger_progress_bar = logging.getLogger("progress_bar")
 
 
 def _calc_hash(existing_file: Path):
@@ -45,28 +43,16 @@ def __init__(self, args: Configuration):
             self.master_hash_list = self.scan_existing_files(self.download_directory)
 
     def download(self):
-        if self.args.progress_bar:
-            with logging_redirect_tqdm():
-                for generator in tqdm(self.reddit_lists, desc="Subreddits", unit="subreddit", colour="red"):
-                    desc = generator.url if isinstance(generator, praw.models.ListingGenerator) else "Posts"
-                    for submission in tqdm(list(generator), desc=desc, unit="post", colour="green", leave=False):
-                        title_short = submission.title[:60] + (submission.title[60:] and "...")
-                        log_str = f"{submission.score:5d}🔼 {title_short}"
-                        try:
-                            if self._download_submission(submission):
-                                logger_progress_bar.info(f"✅ {log_str}")
-                            else:
-                                logger_progress_bar.info(f"❌ {log_str}")
-                        except prawcore.PrawcoreException as e:
-                            logger.error(f"Submission {submission.id} failed to download due to a PRAW exception: {e}")
-                            logger_progress_bar.info(f"❌ {log_str}")
-        else:
-            for generator in self.reddit_lists:
-                for submission in generator:
+        progress = Progress(self.args.progress_bar)
+        with progress.context():
+            for generator in progress.wrap_subreddit_iter(self.reddit_lists):
+                for submission in progress.wrap_post_iter(generator):
                     try:
-                        self._download_submission(submission)
+                        status = self._download_submission(submission)
+                        progress.post_log(submission, status)
                     except prawcore.PrawcoreException as e:
                         logger.error(f"Submission {submission.id} failed to download due to a PRAW exception: {e}")
+                        progress.post_log(submission, False)
 
     def _download_submission(self, submission: praw.models.Submission):
         if submission.id in self.excluded_submission_ids:

bdfr/progress_bar.py

@@ -0,0 +1,41 @@
+import logging
+from contextlib import nullcontext
+
+from tqdm import tqdm
+from tqdm.contrib.logging import logging_redirect_tqdm
+
+logger_progress_bar = logging.getLogger("progress_bar")
+
+
+class Progress:
+    def __init__(self, progress_bar: bool):
+        self.progress_bar = progress_bar
+
+    def context(self):
+        if self.progress_bar:
+            return logging_redirect_tqdm()
+        else:
+            return nullcontext()
+
+    def wrap_subreddit_iter(self, iterator):
+        if self.progress_bar:
+            return tqdm(iterator, desc="Subreddits", unit="subreddit", colour="red")
+        else:
+            return iterator
+
+    def wrap_post_iter(self, generator):
+        if self.progress_bar:
+            try:
+                desc = generator.url
+            except:
+                desc = "Posts"
+            return tqdm(list(generator), desc=desc, unit="post", colour="green", leave=False)
+        else:
+            return generator
+
+    def post_log(self, submission, success: bool):
+        if self.progress_bar:
+            title_short = submission.title[:60] + (submission.title[60:] and "...")
+            log_str = f"{submission.score:5d}🔼 {title_short}"
+            icon = "✅" if success else "❌"
+            logger_progress_bar.info(f"{icon} {log_str}")