From 1511fd6dcc522c2cfbb33ff5daf08dff7b7a4617 Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Sat, 20 Jul 2019 13:44:08 +0100
Subject: [PATCH] add preliminary tqdm work

---
 dvc/progress.py | 136 ++++++++++++++----------------------------------
 setup.py        |   1 +
 2 files changed, 39 insertions(+), 98 deletions(-)

diff --git a/dvc/progress.py b/dvc/progress.py
index 31c882f52..81d4ba658 100644
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -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
diff --git a/setup.py b/setup.py
index 9e98ace7a..0e74366af 100644
--- a/setup.py
+++ b/setup.py
@@ -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'",
 ]