import os import struct import sys import time import typing from typing import Union try: import fcntl import termios import signal _CAN_RESIZE_TERMINAL = True except ImportError: _CAN_RESIZE_TERMINAL = False from . import formatSize, formatTime class AnsiCodes(object): _USE_ANSI = False _CSI = '\033[' def __init__(self, codes): def code_to_chars(code): return AnsiCodes._CSI + str(code) + 'm' for name in dir(codes): if not name.startswith('_'): value = getattr(codes, name) setattr(self, name, code_to_chars(value)) # Add color function for reset_name in ("RESET_%s" % name, "RESET"): if hasattr(codes, reset_name): reset_value = getattr(codes, reset_name) setattr(self, "%s" % name.lower(), AnsiCodes._mkfunc(code_to_chars(value), code_to_chars(reset_value))) break @staticmethod def _mkfunc(color, reset): def _cwrap(text, *styles): if not AnsiCodes._USE_ANSI: return text s = '' for st in styles: s += st s += color + text + reset if styles: s += Style.RESET_ALL return s return _cwrap def __getattribute__(self, name): attr = super(AnsiCodes, self).__getattribute__(name) if (hasattr(attr, "startswith") and attr.startswith(AnsiCodes._CSI) and not AnsiCodes._USE_ANSI): return "" else: return attr def __getitem__(self, name): return getattr(self, name.upper()) @classmethod def init(cls, allow_colors): cls._USE_ANSI = allow_colors and cls._term_supports_color() @staticmethod def _term_supports_color(): if (os.environ.get("TERM") == "dumb" or os.environ.get("OS") == "Windows_NT"): return False return hasattr(sys.stdout, "isatty") and sys.stdout.isatty() class AnsiFore: GREY = 30 # noqa RED = 31 # noqa GREEN = 32 # noqa YELLOW = 33 # noqa BLUE = 34 # noqa MAGENTA = 35 # noqa CYAN = 36 # noqa WHITE = 37 # noqa RESET = 39 # noqa class AnsiBack: GREY = 40 # noqa RED = 41 # noqa GREEN = 42 # noqa YELLOW = 43 # noqa BLUE = 44 # noqa MAGENTA = 45 # noqa CYAN = 46 # noqa WHITE = 47 # noqa RESET = 49 # noqa class AnsiStyle: RESET_ALL = 0 # noqa BRIGHT = 1 # noqa RESET_BRIGHT = 22 # noqa DIM = 2 # noqa RESET_DIM = RESET_BRIGHT # noqa ITALICS = 3 # noqa RESET_ITALICS = 23 # noqa UNDERLINE = 4 # noqa RESET_UNDERLINE = 24 # noqa BLINK_SLOW = 5 # noqa RESET_BLINK_SLOW = 25 # noqa BLINK_FAST = 6 # noqa RESET_BLINK_FAST = 26 # noqa INVERSE = 7 # noqa RESET_INVERSE = 27 # noqa STRIKE_THRU = 9 # noqa RESET_STRIKE_THRU = 29 # noqa Fore = AnsiCodes(AnsiFore) Back = AnsiCodes(AnsiBack) Style = AnsiCodes(AnsiStyle) def ERROR_COLOR(): return Fore.RED def WARNING_COLOR(): return Fore.YELLOW def HEADER_COLOR(): return Fore.GREEN class Spinner(object): """ A class to display a spinner in the terminal. It is designed to be used with the `with` statement:: with Spinner("Reticulating splines", "green") as s: for item in enumerate(items): s.next() """ _default_unicode_chars = "◓◑◒◐" _default_ascii_chars = "-/|\\" def __init__(self, msg, file=None, step=1, chars=None, use_unicode=True, print_done=True): self._msg = msg self._file = file or sys.stdout self._step = step if not chars: if use_unicode: chars = self._default_unicode_chars else: chars = self._default_ascii_chars self._chars = chars self._silent = not self._file.isatty() self._print_done = print_done def _iterator(self): chars = self._chars index = 0 write = self._file.write flush = self._file.flush while True: write("\r") write(self._msg) write(" ") write(chars[index]) flush() yield for _ in range(self._step): yield index += 1 if index == len(chars): index = 0 def __enter__(self): if self._silent: return self._silent_iterator() else: return self._iterator() def __exit__(self, exc_type, exc_value, traceback): write = self._file.write flush = self._file.flush if not self._silent: write("\r") write(self._msg) if self._print_done: if exc_type is None: write(Fore.GREEN + ' [Done]\n') else: write(Fore.RED + ' [Failed]\n') else: write("\n") flush() def _silent_iterator(self): self._file.write(self._msg) self._file.flush() while True: yield class ProgressBar(object): """ A class to display a progress bar in the terminal. It is designed to be used either with the `with` statement:: with ProgressBar(len(items)) as bar: for item in enumerate(items): bar.update() or as a generator:: for item in ProgressBar(items): item.process() """ def __init__(self, total_or_items: Union[int, typing.Sequence], file=None): """ total_or_items : int or sequence If an int, the number of increments in the process being tracked. If a sequence, the items to iterate over. file : writable file-like object, optional The file to write the progress bar to. Defaults to `sys.stdout`. If `file` is not a tty (as determined by calling its `isatty` member, if any), the scrollbar will be completely silent. """ self._file = file or sys.stdout if not self._file.isatty(): self.update = self._silent_update self._silent = True else: self._silent = False try: self._items = iter(total_or_items) self._total = len(total_or_items) except TypeError: try: self._total = int(total_or_items) self._items = iter(range(self._total)) except TypeError: raise TypeError("First argument must be int or sequence") self._start_time = time.time() self._should_handle_resize = ( _CAN_RESIZE_TERMINAL and self._file.isatty()) self._handle_resize() if self._should_handle_resize: signal.signal(signal.SIGWINCH, self._handle_resize) self._signal_set = True else: self._signal_set = False self.update(0) def _handle_resize(self, signum=None, frame=None): self._terminal_width = getTtySize(self._file, self._should_handle_resize)[1] def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): if not self._silent: if exc_type is None: self.update(self._total) self._file.write('\n') self._file.flush() if self._signal_set: signal.signal(signal.SIGWINCH, signal.SIG_DFL) def __iter__(self): return self def __next__(self): try: rv = next(self._items) except StopIteration: self.__exit__(None, None, None) raise else: self.update() return rv def update(self, value=None): """ Update the progress bar to the given value (out of the total given to the constructor). """ if value is None: value = self._current_value = self._current_value + 1 else: self._current_value = value if self._total == 0: frac = 1.0 else: frac = float(value) / float(self._total) file = self._file write = file.write suffix = self._formatSuffix(value, frac) self._bar_length = self._terminal_width - 37 bar_fill = int(float(self._bar_length) * frac) write("\r|") write(Fore.BLUE + '=' * bar_fill + Fore.RESET) if bar_fill < self._bar_length: write(Fore.GREEN + '>' + Fore.RESET) write("-" * (self._bar_length - bar_fill - 1)) write("|") write(suffix) self._file.flush() def _formatSuffix(self, value, frac): if value >= self._total: t = time.time() - self._start_time time_str = ' ' elif value <= 0: t = None time_str = '' else: t = ((time.time() - self._start_time) * (1.0 - frac)) / frac time_str = ' ETA ' if t is not None: time_str += formatTime(t, short=True) suffix = ' {0:>4s}/{1:>4s}'.format(formatSize(value, short=True), formatSize(self._total, short=True)) suffix += ' ({0:>6s}%)'.format("{0:.2f}".format(frac * 100.0)) suffix += time_str return suffix def _silent_update(self, value=None): pass @classmethod def map(cls, function, items, multiprocess=False, file=None): """ Does a `map` operation while displaying a progress bar with percentage complete. :: def work(i): print(i) ProgressBar.map(work, range(50)) Parameters: function : function Function to call for each step items : sequence Sequence where each element is a tuple of arguments to pass to *function*. multiprocess : bool, optional If `True`, use the `multiprocessing` module to distribute each task to a different processor core. file : writeable file-like object, optional The file to write the progress bar to. Defaults to `sys.stdout`. If `file` is not a tty (as determined by calling its `isatty` member, if any), the scrollbar will be completely silent. """ results = [] if file is None: file = sys.stdout with cls(len(items), file=file) as bar: step_size = max(200, bar._bar_length) steps = max(int(float(len(items)) / step_size), 1) if not multiprocess: for i, item in enumerate(items): function(item) if (i % steps) == 0: bar.update(i) else: import multiprocessing p = multiprocessing.Pool() for i, result in enumerate(p.imap_unordered(function, items, steps)): bar.update(i) results.append(result) return results def printMsg(s): fp = sys.stdout assert isinstance(s, str) try: fp.write("%s\n" % s) except UnicodeEncodeError: fp.write("%s\n" % str(s.encode("utf-8", "replace"), "utf-8")) fp.flush() def printError(s): _printWithColor(s, ERROR_COLOR(), sys.stderr) def printWarning(s): _printWithColor(s, WARNING_COLOR(), sys.stdout) def printHeader(s): _printWithColor(s, HEADER_COLOR(), sys.stdout) def boldText(s, c=None): return formatText(s, b=True, c=c) def formatText(s, b=False, c=None): return ((Style.BRIGHT if b else '') + (c or '') + s + (Fore.RESET if c else '') + (Style.RESET_BRIGHT if b else '')) def _printWithColor(s, color, file): assert isinstance(s, str) file.write(color + s + Fore.RESET + '\n') file.flush() def cformat(msg, fg, bg=None, styles=None): """Format ``msg`` with foreground and optional background. Optional ``styles`` lists will also be applied. The formatted string is returned.""" fg = fg or "" bg = bg or "" styles = "".join(styles or []) reset = Fore.RESET + Back.RESET + Style.RESET_ALL if (fg or bg or styles) \ else "" output = "%(fg)s%(bg)s%(styles)s%(msg)s%(reset)s" % locals() return output def getTtySize(fd=sys.stdout, check_tty=True): hw = None if check_tty: try: data = fcntl.ioctl(fd, termios.TIOCGWINSZ, '\0' * 4) hw = struct.unpack("hh", data) except (OSError, NameError): pass if not hw: try: hw = (int(os.environ.get('LINES')), int(os.environ.get('COLUMNS'))) except (TypeError, ValueError): hw = (78, 25) return hw def cprint(msg, fg, bg=None, styles=None, file=sys.stdout): """Calls ``cformat`` and prints the result to output stream ``file``.""" print(cformat(msg, fg, bg=bg, styles=styles), file=file) if __name__ == "__main__": AnsiCodes.init(True) def checkCode(c): return (c[0] != '_' and "RESET" not in c and c[0] == c[0].upper() ) for bg_name, bg_code in ((c, getattr(Back, c)) for c in dir(Back) if checkCode(c)): sys.stdout.write('%s%-7s%s %s ' % (bg_code, bg_name, Back.RESET, bg_code)) for _, fg_code in ((c, getattr(Fore, c)) for c in dir(Fore) if checkCode(c)): sys.stdout.write(fg_code) for st_name, st_code in ((c, getattr(Style, c)) for c in dir(Style) if checkCode(c)): sys.stdout.write('%s%s %s %s' % (st_code, st_name, getattr(Style, "RESET_%s" % st_name), bg_code)) sys.stdout.write("%s\n" % Style.RESET_ALL) sys.stdout.write("\n") with Spinner(Fore.GREEN + "Phase #1") as spinner: for _ in range(50): time.sleep(.05) next(spinner) with Spinner(Fore.RED + "Phase #2" + Fore.RESET, print_done=False) as spinner: for _ in range(50): time.sleep(.05) next(spinner) with Spinner("Phase #3", print_done=False, use_unicode=False) as spinner: for _ in range(50): next(spinner) time.sleep(.05) with Spinner("Phase #4", print_done=False, chars='.oO°Oo.') as spinner: for _ in range(50): next(spinner) time.sleep(.05) items = [x for x in range(200)] with ProgressBar(len(items)) as bar: for _ in enumerate(items): bar.update() time.sleep(.05) for _ in iter(ProgressBar(items)): time.sleep(.05) progress = 0 max_val = 320000000 with ProgressBar(max_val) as bar: while progress < max_val: progress += 23400 bar.update(progress) time.sleep(.001)