add preliminary tqdm work

This commit is contained in:
Casper da Costa-Luis
2019-08-10 14:11:02 +01:00
parent 1b2823b5c5
commit 1511fd6dcc
2 changed files with 39 additions and 98 deletions
+38 -98
View File
@@ -2,129 +2,68 @@
from __future__ import print_function
from __future__ import unicode_literals
import logging
from tqdm import tqdm
from dvc.utils.compat import str
import sys
import threading
CLEARLINE_PATTERN = "\r\x1b[K"
class Progress(object):
class Progress(tqdm):
"""
Simple multi-target progress bar.
"""
def __init__(self):
self._n_total = 0
self._n_finished = 0
self._lock = threading.Lock()
self._line = None
super(Progress, self).__init__(
total=0,
disable=logging.getLogger(__name__).getEffectiveLevel() >= logging.CRITICAL)
self.set_lock(threading.Lock())
self._targets = {}
self.clearln()
def set_n_total(self, total):
"""Sets total number of targets."""
self._n_total = total
self._n_finished = 0
self.reset(total)
@property
def is_finished(self):
"""Returns if all targets have finished."""
return self._n_total == self._n_finished
return self.total == self.n
def clearln(self):
self._print(CLEARLINE_PATTERN, end="")
self.display("")
def _writeln(self, line):
self.clearln()
self._print(line, end="")
sys.stdout.flush()
def reset(self):
with self._lock:
self._n_total = 0
self._n_finished = 0
self._line = None
def refresh(self, line=None):
"""Refreshes progress bar."""
# Just go away if it is locked. Will update next time
if not self._lock.acquire(False):
return
if line is None:
line = self._line
if sys.stdout.isatty() and line is not None:
self._writeln(line)
self._line = line
self._lock.release()
def update_target(self, name, current, total):
def update_target(self, name, current, total, auto_finish=False):
"""Updates progress bar for a specified target."""
self.refresh(self._bar(name, current, total))
if name in self._targets:
t = self._targets[name]
else:
with self.get_lock():
# TODO: up to 10 nested bars
position = 1 + self._get_free_pos(self) % 10
t = tqdm(
total=total,
desc=name,
leave=False,
position=position)
self._targets[name] = t
self.total += 1
t.update(current - t.n)
if auto_finish and t.n == t.total:
self.finish_target(name)
def finish_target(self, name):
"""Finishes progress bar for a specified target."""
# We have to write a msg about finished target
with self._lock:
pbar = self._bar(name, 100, 100)
if sys.stdout.isatty():
self.clearln()
self._print(pbar)
self._n_finished += 1
self._line = None
def _bar(self, target_name, current, total):
"""
Make a progress bar out of info, which looks like:
(1/2): [########################################] 100% master.zip
"""
bar_len = 30
if total is None:
state = 0
percent = "?% "
else:
total = int(total)
state = int((100 * current) / total) if current < total else 100
percent = str(state) + "% "
if self._n_total > 1:
num = "({}/{}): ".format(self._n_finished + 1, self._n_total)
else:
num = ""
n_sh = int((state * bar_len) / 100)
n_sp = bar_len - n_sh
pbar = "[" + "#" * n_sh + " " * n_sp + "] "
return num + pbar + percent + target_name
@staticmethod
def _print(*args, **kwargs):
import logging
logger = logging.getLogger(__name__)
if logger.getEffectiveLevel() == logging.CRITICAL:
return
print(*args, **kwargs)
def __enter__(self):
self._lock.acquire(True)
if self._line is not None:
self.clearln()
def __exit__(self, typ, value, tbck):
if self._line is not None:
self.refresh()
self._lock.release()
t = self._targets.pop(name)
t.close()
# TODO: We have to write a msg about finished target
if self.total < 100:
# only if less that 100 items
print(t)
self.update()
def __call__(self, seq, name="", total=None):
if total is None:
@@ -138,6 +77,7 @@ class Progress(object):
class ProgressCallback(object):
# TODO: is this meant to be a thin wrapper for multi-progress, or just one bar?
def __init__(self, total):
self.total = total
self.current = 0
+1
View File
@@ -61,6 +61,7 @@ install_requires = [
"funcy>=1.12",
"pathspec>=0.5.9",
"shortuuid>=0.5.0",
"tqdm>=4.32.1",
"win-unicode-console>=0.5; sys_platform == 'win32'",
]