diff --git a/main.py b/main.py index 620c348..b8a1a94 100755 --- a/main.py +++ b/main.py @@ -8,6 +8,7 @@ from typing import Optional from logger import color_option, logging, setup_logging, verbose_option from wrapper import nowrapper_option, enforce_wrap +from progressbar import progress_bars_option from config.cli import config, config_option, cmd_config from packages.cli import cmd_packages @@ -25,15 +26,18 @@ from image.cli import cmd_image @config_option @nowrapper_option @color_option +@progress_bars_option def cli( verbose: bool = False, config_file: Optional[str] = None, wrapper_override: Optional[bool] = None, error_shell: bool = False, force_colors: Optional[bool] = None, + force_progress_bars: Optional[bool] = None, ): setup_logging(verbose, force_colors=force_colors) config.runtime.verbose = verbose + config.runtime.progress_bars = force_progress_bars config.runtime.no_wrap = wrapper_override is False config.runtime.error_shell = error_shell config.try_load_file(config_file) diff --git a/progressbar.py b/progressbar.py new file mode 100644 index 0000000..41104ed --- /dev/null +++ b/progressbar.py @@ -0,0 +1,51 @@ +import click +import sys + +from enlighten import Counter, Manager, get_manager as _getmanager +from typing import Hashable, Optional + +from config.state import config + +DEFAULT_OUTPUT = sys.stdout + +managers: dict[Hashable, Manager] = {} + +progress_bars_option = click.option( + '--force-progress-bars/--no-progress-bars', + is_flag=True, + default=None, + help='Force enable/disable progress bars. Defaults to autodetection.', +) + + +def get_manager(file=DEFAULT_OUTPUT, enabled: Optional[bool] = None) -> Manager: + global managers + m = managers.get(file, None) + if not m: + kwargs = {} + if enabled is None or config.runtime.progress_bars is False: + enabled = config.runtime.progress_bars + if enabled is not None: + kwargs = {"enabled": enabled} + m = _getmanager(file, **kwargs) + managers[file] = m + return m + + +def get_progress_bar(*kargs, file=DEFAULT_OUTPUT, leave=False, **kwargs) -> Counter: + m = get_manager(file=file) + + kwargs["file"] = file + kwargs["leave"] = leave + return m.counter(*kargs, **kwargs) + + +def get_levels_bar(*kargs, file=DEFAULT_OUTPUT, enable_rate=True, **kwargs): + kwargs["fields"] = {"name": "None", "level": 1, "levels_total": 1} | (kwargs.get("fields", None) or {}) + f = (u'{desc}: {name}{desc_pad}{percentage:3.0f}%|{bar}| ' + u'{count:{len_total}d}/{total:d} ' + u'[lvl: {level}/{levels_total}] ') + if enable_rate: + f += u'[{elapsed}<{eta}, {rate:.2f}{unit_pad}{unit}/s]' + kwargs["bar_format"] = f + return get_progress_bar(*kargs, **kwargs) diff --git a/requirements.txt b/requirements.txt index 728829b..25e5f4a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ munch setuptools # required by munch requests python-dateutil +enlighten