1"""Manages progress bars for dvc repo.""" 2 3from __future__ import print_function 4from __future__ import unicode_literals 5 6from dvc.utils.compat import str 7 8import sys 9import threading 10 11 12class Progress(object): 13 """ 14 Simple multi-target progress bar. 15 """ 16 17 def __init__(self): 18 self._n_total = 0 19 self._n_finished = 0 20 self._lock = threading.Lock() 21 self._line = None 22 23 def set_n_total(self, total): 24 """Sets total number of targets.""" 25 self._n_total = total 26 self._n_finished = 0 27 28 @property 29 def is_finished(self): 30 """Returns if all targets have finished.""" 31 return self._n_total == self._n_finished 32 33 def _clearln(self): 34 self._print("\r\x1b[K", end="") 35 36 def _writeln(self, line): 37 self._clearln() 38 self._print(line, end="") 39 sys.stdout.flush() 40 41 def refresh(self, line=None): 42 """Refreshes progress bar.""" 43 # Just go away if it is locked. Will update next time 44 if not self._lock.acquire(False): 45 return 46 47 if line is None: 48 line = self._line 49 50 if sys.stdout.isatty() and line is not None: 51 self._writeln(line) 52 self._line = line 53 54 self._lock.release() 55 56 def update_target(self, name, current, total): 57 """Updates progress bar for a specified target.""" 58 self.refresh(self._bar(name, current, total)) 59 60 def finish_target(self, name): 61 """Finishes progress bar for a specified target.""" 62 # We have to write a msg about finished target 63 with self._lock: 64 pbar = self._bar(name, 100, 100) 65 66 if sys.stdout.isatty(): 67 self._clearln() 68 69 self._print(pbar) 70 71 self._n_finished += 1 72 self._line = None 73 74 def _bar(self, target_name, current, total): 75 """ 76 Make a progress bar out of info, which looks like: 77 (1/2): [########################################] 100% master.zip 78 """ 79 bar_len = 30 80 81 if total is None: 82 state = 0 83 percent = "?% " 84 else: 85 total = int(total) 86 state = int((100 * current) / total) if current < total else 100 87 percent = str(state) + "% " 88 89 if self._n_total > 1: 90 num = "({}/{}): ".format(self._n_finished + 1, self._n_total) 91 else: 92 num = "" 93 94 n_sh = int((state * bar_len) / 100) 95 n_sp = bar_len - n_sh 96 pbar = "[" + "#" * n_sh + " " * n_sp + "] " 97 98 return num + pbar + percent + target_name 99 100 @staticmethod 101 def _print(*args, **kwargs): 102 import dvc.logger as logger 103 104 if logger.is_quiet(): 105 return 106 107 print(*args, **kwargs) 108 109 def __enter__(self): 110 self._lock.acquire(True) 111 if self._line is not None: 112 self._clearln() 113 114 def __exit__(self, typ, value, tbck): 115 if self._line is not None: 116 self.refresh() 117 self._lock.release() 118 119 120def progress_aware(f): 121 """ Decorator to add a new line if progress bar hasn't finished """ 122 from functools import wraps 123 124 @wraps(f) 125 def wrapper(*args, **kwargs): 126 if not progress.is_finished: 127 progress._print() 128 129 return f(*args, **kwargs) 130 131 return wrapper 132 133 134progress = Progress() # pylint: disable=invalid-name 135