diff --git a/pytchat/cli/__init__.py b/pytchat/cli/__init__.py index c11d1b5..ff04ec0 100644 --- a/pytchat/cli/__init__.py +++ b/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,15 +71,18 @@ def main(): print(f" output path: {path.resolve()}") duration = info.get_duration() - pbar = ProgressBar(total=(duration * 1000) / 0.99, status="Extracting") - ex = Extractor(video_id, + pbar = ProgressBar(total=(duration * 1000), status="Extracting") + ex = Extractor(video_id, callback=pbar._disp, div=10) signal.signal(signal.SIGINT, (lambda a, b: cancel(ex, pbar))) data = ex.extract() if data == []: return False - pbar.reset("#", "=", total=len(data), status="Rendering ") + 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() - pbar.reset('#', '#', status='Completed ') - pbar.close() + 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() diff --git a/pytchat/cli/arguments.py b/pytchat/cli/arguments.py index 0e8e9e8..e22b293 100644 --- a/pytchat/cli/arguments.py +++ b/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 diff --git a/pytchat/cli/progressbar_rich.py b/pytchat/cli/progressbar_rich.py new file mode 100644 index 0000000..65ca9b0 --- /dev/null +++ b/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 diff --git a/pytchat/cli/progressbar_simple.py b/pytchat/cli/progressbar_simple.py new file mode 100644 index 0000000..776fad2 --- /dev/null +++ b/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