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()