Add progress bar

pytchat/cli/__init__.py

@@ -1,12 +1,17 @@
 import argparse
+
 import os
+import signal
+from json.decoder import JSONDecodeError
 from pathlib import Path
-from pytchat.util.extract_video_id import extract_video_id
 from .arguments import Arguments
-from .. exceptions import InvalidVideoIdException, NoContents, VideoInfoParseException
+from .progressbar import ProgressBar
+from .. exceptions import InvalidVideoIdException, NoContents, PatternUnmatchError
 from .. processors.html_archiver import HTMLArchiver
 from .. tool.extract.extractor import Extractor
 from .. tool.videoinfo import VideoInfo
+from .. util.extract_video_id import extract_video_id
+from .. import util
 from .. import __version__
 
 '''
@@ -35,40 +40,57 @@ def main():
         return
 
     # Extractor
-    if Arguments().video_ids:
-        for video_id in Arguments().video_ids:
-            if '[' in video_id:
-                video_id = video_id.replace('[', '').replace(']', '')
-            try:
-                video_id = extract_video_id(video_id)
-                if os.path.exists(Arguments().output):
-                    path = Path(Arguments().output + video_id + '.html')
-                else:
-                    raise FileNotFoundError
-                info = VideoInfo(video_id)
-                print(f"Extracting...\n"
-                      f" video_id: {video_id}\n"
-                      f" channel:  {info.get_channel_name()}\n"
-                      f" title:    {info.get_title()}")
-
-                print(f" output path: {path.resolve()}")
-                Extractor(video_id,
-                          processor=HTMLArchiver(
-                              Arguments().output + video_id + '.html'),
-                          callback=_disp_progress
-                          ).extract()
-                print("\nExtraction end.\n")
-            except InvalidVideoIdException:
-                print("Invalid Video ID or URL:", video_id)
-            except (TypeError, NoContents) as e:
-                print(e)
-            except FileNotFoundError:
-                print("The specified directory does not exist.:{}".format(Arguments().output))
-            except VideoInfoParseException:
-                print("Cannot parse video information.:{}".format(video_id))
+    if not Arguments().video_ids:
+        parser.print_help()
         return
-    parser.print_help()
+    for video_id in Arguments().video_ids:
+        if '[' in video_id:
+            video_id = video_id.replace('[', '').replace(']', '')
+        try:
+            video_id = extract_video_id(video_id)
+            if os.path.exists(Arguments().output):
+                path = Path(Arguments().output + video_id + '.html')
+            else:
+                raise FileNotFoundError
+            info = VideoInfo(video_id)
+            print(f"Extracting...\n"
+                  f" video_id: {video_id}\n"
+                  f" channel:  {info.get_channel_name()}\n"
+                  f" title:    {info.get_title()}")
+
+            print(f" output path: {path.resolve()}")
+            duration = info.get_duration()
+            pbar = ProgressBar(duration)
+            ex = Extractor(video_id,
+                    processor=HTMLArchiver(Arguments().output + video_id + '.html'),
+                    callback=pbar._disp,
+                    div=10)
+            signal.signal(signal.SIGINT, (lambda a, b: cancel(ex, pbar)))
+            ex.extract()
+            pbar.close()
+            if pbar.is_cancelled():
+                print("\nThe extraction process has been discontinued.\n")
+                return
+            print("\nThe extraction process has been completed.\n")
+        except InvalidVideoIdException:
+            print("Invalid Video ID or URL:", video_id)
+        except (TypeError, NoContents) as e:
+
+            print(e.with_traceback())
+        except FileNotFoundError:
+            print("The specified directory does not exist.:{}".format(Arguments().output))
+        except JSONDecodeError as e:
+            print(e.msg)
+            print("Cannot parse video information.:{}".format(video_id))
+            util.save(e.doc, "ERR_JSON_DECODE", ".dat")
+        except PatternUnmatchError as e:
+            print(e.msg)
+            print("Cannot parse video information.:{}".format(video_id))
+            util.save(e.doc, "ERR_PATTERN_UNMATCH", ".dat")
+
+    return
 
 
-def _disp_progress(a, b):
-    print('.', end="", flush=True)
+def cancel(ex: Extractor, pbar: ProgressBar):
+    ex.cancel()
+    pbar.cancel()

pytchat/cli/progressbar.py

@@ -0,0 +1,41 @@
+'''
+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, duration):
+        self._duration = duration
+        self._count = 0
+        self._bar_len = 60
+        self._cancelled = False
+
+    def _disp(self, _, fetched):
+        self._progress(fetched / 1000, self._duration)
+
+    def _progress(self, fillin, total, status=''):
+        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 filled_len > self._bar_len:
+            filled_len = self._bar_len
+            percents = 100
+        bar = '=' * filled_len + ' ' * (self._bar_len - filled_len)
+        sys.stdout.write(' [%s] %s%s ...%s\r' % (bar, percents, '%', status))
+        sys.stdout.flush()
+
+    def close(self):
+        if not self._cancelled:
+            self._progress(self._duration, self._duration)
+
+    def cancel(self):
+        self._cancelled = True
+    
+    def is_cancelled(self):
+        return self._cancelled