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