Make it possible to switch progress bar

pytchat/cli/__init__.py

@@ -1,11 +1,11 @@
 import argparse
 
 import os
+import sys
 import signal
 from json.decoder import JSONDecodeError
 from pathlib import Path
 from .arguments import Arguments
-from .progressbar import ProgressBar
 from .. exceptions import InvalidVideoIdException, NoContents, PatternUnmatchError
 from .. processors.html_archiver import HTMLArchiver
 from .. tool.extract.extractor import Extractor
@@ -24,6 +24,8 @@ https://github.com/PetterKraabol/Twitch-Chat-Downloader
 
 
 def main():
+    
+
     # Arguments
     parser = argparse.ArgumentParser(description=f'pytchat v{__version__}')
     parser.add_argument('-v', f'--{Arguments.Name.VIDEO_IDS}', type=str,
@@ -32,11 +34,18 @@ def main():
                         'If ID starts with a hyphen (-), enclose the ID in square brackets.')
     parser.add_argument('-o', f'--{Arguments.Name.OUTPUT}', type=str,
                         help='Output directory (end with "/"). default="./"', default='./')
-    parser.add_argument(f'--{Arguments.Name.VERSION}', action='store_true',
-                        help='Show version')
+    parser.add_argument(f'--{Arguments.Name.PBAR}', action='store_true',
+                        help='Display rich progress bar')
     parser.add_argument(f'--{Arguments.Name.SAVE_ERROR_DATA}', action='store_true',
                         help='Save error data when error occurs(".dat" file)')
+    parser.add_argument(f'--{Arguments.Name.VERSION}', action='store_true',
+                        help='Show version')
     Arguments(parser.parse_args().__dict__)
+
+    if Arguments().pbar:
+        from .progressbar_rich import ProgressBar
+    else:
+        from .progressbar_simple import ProgressBar
     if Arguments().print_version:
         print(f'pytchat v{__version__}     © 2019 taizan-hokuto')
         return
@@ -62,7 +71,7 @@ def main():
 
             print(f" output path: {path.resolve()}")
             duration = info.get_duration()
-            pbar = ProgressBar(total=(duration * 1000) / 0.99, status="Extracting")
+            pbar = ProgressBar(total=(duration * 1000), status="Extracting")
             ex = Extractor(video_id,                  
                     callback=pbar._disp,
                     div=10)
@@ -70,7 +79,10 @@ def main():
             data = ex.extract()
             if data == []:
                 return False
+            if Arguments().pbar:
                 pbar.reset("#", "=", total=len(data), status="Rendering  ")
+            else:
+                pbar.reset("=", "", total=len(data), status="Rendering  ")
             processor = HTMLArchiver(Arguments().output + video_id + '.html', callback=pbar._disp)
             processor.process(
                 [{'video_id': None,
@@ -78,8 +90,13 @@ def main():
                 'chatdata': (action["replayChatItemAction"]["actions"][0] for action in data)}]
             )
             processor.finalize()
+            if Arguments().pbar:
                 pbar.reset('#', '#', status='Completed   ')
                 pbar.close()
+            else:
+                pbar.close()
+                print("\nCompleted")
+            
             print()
             if pbar.is_cancelled():
                 print("\nThe extraction process has been discontinued.\n")
@@ -106,6 +123,6 @@ def main():
     return
 
 
-def cancel(ex: Extractor, pbar: ProgressBar):
+def cancel(ex, pbar):
     ex.cancel()
     pbar.cancel()

pytchat/cli/arguments.py

@@ -19,6 +19,7 @@ class Arguments(metaclass=Singleton):
         OUTPUT: str = 'output_dir'
         VIDEO_IDS: str = 'video_id'
         SAVE_ERROR_DATA: bool = 'save_error_data'
+        PBAR: bool ='pbar'
 
     def __init__(self,
                  arguments: Optional[Dict[str, Union[str, bool, int]]] = None):
@@ -36,6 +37,7 @@ class Arguments(metaclass=Singleton):
         self.output: str = arguments[Arguments.Name.OUTPUT]
         self.video_ids: List[int] = []
         self.save_error_data: bool = arguments[Arguments.Name.SAVE_ERROR_DATA]
+        self.pbar: bool = arguments[Arguments.Name.PBAR]
         # Videos
         if arguments[Arguments.Name.VIDEO_IDS]:
             self.video_ids = [video_id

pytchat/cli/progressbar_rich.py

@@ -0,0 +1,52 @@
+'''
+This code for this progress bar is based on
+vladignatyev/progress.py
+https://gist.github.com/vladignatyev/06860ec2040cb497f0f3
+(MIT License)
+'''
+import sys
+
+
+class ProgressBar:
+    def __init__(self, total, status):
+        self._bar_len = 60
+        self._cancelled = False
+        self.reset(total=total, status=status)
+        self._blinker = 0
+
+    def reset(self, symbol_done="=", symbol_space=" ", total=100, status=''):
+        self._symbol_done = symbol_done
+        self._symbol_space = symbol_space
+        self._total = total
+        self._status = status
+        self._count = 0
+
+    def _disp(self, _, fetched):
+        self._progress(fetched, self._total)
+
+    def _progress(self, fillin, total):
+        if total == 0 or self._cancelled:
+            return
+        self._count += fillin
+        filled_len = int(round(self._bar_len * self._count / float(total)))
+        percents = round(100.0 * self._count / float(total), 1)
+        if percents > 100:
+            percents = 100.0
+        if filled_len > self._bar_len:
+            filled_len = self._bar_len
+            
+        bar = self._symbol_done * filled_len + \
+              self._symbol_space * (self._bar_len - filled_len)
+        sys.stdout.write(' [%s] %s%s ...%s \r' % (bar, percents, '%', self._status))
+        sys.stdout.flush()
+        self._blinker += 1
+
+    def close(self):
+        if not self._cancelled:
+            self._progress(self._total, self._total)
+
+    def cancel(self):
+        self._cancelled = True
+    
+    def is_cancelled(self):
+        return self._cancelled

pytchat/cli/progressbar_simple.py

@@ -0,0 +1,49 @@
+'''
+This code for this progress bar is based on
+vladignatyev/progress.py
+https://gist.github.com/vladignatyev/06860ec2040cb497f0f3
+(MIT License)
+'''
+import sys
+
+
+class ProgressBar:
+    def __init__(self, total, status):
+        self._bar_len = 60
+        self._cancelled = False
+        print(''.join([' ' * 10, '|', '-' * (self._bar_len), '|']), end="")
+        self.reset(total=total, status=status)
+
+    def reset(self, symbol_done="=", symbol_space=" ", total=100, status=''):
+        self._symbol_done = symbol_done
+        self._symbol_space = symbol_space
+        self._total = total
+        self._status = status
+        self._old_len = 0
+        self._count = 0
+        print()
+        print(f'{status:<11}', end='')
+
+    def _disp(self, _, fetched):
+        self._progress(fetched, self._total)
+
+    def _progress(self, fillin, total):
+        if total == 0 or self._cancelled:
+            return
+        self._count += fillin
+        filled_len = int(round(self._bar_len * self._count / float(total)))
+        if filled_len > self._bar_len:
+            filled_len = self._bar_len
+        print((filled_len - self._old_len) * self._symbol_done, end="")
+        sys.stdout.flush()
+        self._old_len = filled_len
+
+    def close(self):
+        if not self._cancelled:
+            self._progress(self._total, self._total)
+
+    def cancel(self):
+        self._cancelled = True
+    
+    def is_cancelled(self):
+        return self._cancelled