forked from iterative/dvc
-
Notifications
You must be signed in to change notification settings - Fork 0
/
progress.py
110 lines (100 loc) 路 3.43 KB
/
progress.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
"""Manages progress bars for dvc repo."""
from __future__ import print_function
import logging
from tqdm import tqdm
from copy import deepcopy
from concurrent.futures import ThreadPoolExecutor
class TqdmThreadPoolExecutor(ThreadPoolExecutor):
"""
Ensure worker progressbars are cleared away properly.
"""
def __enter__(self):
"""
Creates a blank initial dummy progress bar if needed so that workers
are forced to create "nested" bars.
"""
blank_bar = Tqdm(bar_format="Multi-Threaded:", leave=False)
if blank_bar.pos > 0:
# already nested - don't need a placeholder bar
blank_bar.close()
self.bar = blank_bar
return super(TqdmThreadPoolExecutor, self).__enter__()
def __exit__(self, *a, **k):
super(TqdmThreadPoolExecutor, self).__exit__(*a, **k)
self.bar.close()
class Tqdm(tqdm):
"""
maximum-compatibility tqdm-based progressbars
"""
BAR_FMT_DEFAULT = (
"{percentage:3.0f}%|{bar:10}|{desc} {bar:-10b}{n}/{total}"
" [{elapsed}<{remaining}, {rate_fmt:>11}{postfix}]"
)
BAR_FMT_NOTOTAL = "{desc} {bar:b}{n} [{elapsed}<??:??, {rate_fmt:>11}{postfix}]"
def __init__(
self,
iterable=None,
disable=None,
bytes=False, # pylint: disable=W0622
desc_truncate=None,
leave=None,
bar_format=None,
**kwargs
):
"""
bytes : shortcut for
`unit='B', unit_scale=True, unit_divisor=1024, miniters=1`
desc_truncate : like `desc` but will truncate to 10 chars
kwargs : anything accepted by `tqdm.tqdm()`
"""
kwargs = deepcopy(kwargs)
kwargs.setdefault("unit_scale", True)
if bytes:
for k, v in dict(
unit="B", unit_scale=True, unit_divisor=1024, miniters=1
).items():
kwargs.setdefault(k, v)
if desc_truncate is not None:
kwargs.setdefault("desc", self.truncate(desc_truncate))
if disable is None:
disable = (
logging.getLogger(__name__).getEffectiveLevel()
>= logging.CRITICAL
)
if bar_format is None:
if kwargs.get('total', getattr(iterable, '__len__', None)) is None:
bar_format = self.BAR_FMT_NOTOTAL
else:
bar_format = self.BAR_FMT_DEFAULT
super(Tqdm, self).__init__(
iterable=iterable,
disable=disable,
leave=leave,
bar_format=bar_format,
**kwargs
)
def update_desc(self, desc, n=1, truncate=True):
"""
Calls `set_description_str(truncate(desc))` and `update(n)`
"""
self.set_description_str(
self.truncate(desc) if truncate else desc, refresh=False
)
self.update(n)
def update_to(self, current, total=None):
if total:
self.total = total # pylint: disable=W0613,W0201
self.update(current - self.n)
@classmethod
def truncate(cls, s, max_len=25, end=True, fill="..."):
"""
Guarantee len(output) < max_lenself.
>>> truncate("hello", 4)
'...o'
"""
if len(s) <= max_len:
return s
if len(fill) > max_len:
return fill[-max_len:] if end else fill[:max_len]
i = max_len - len(fill)
return (fill + s[-i:]) if end else (s[:i] + fill)