From 0ebce9c1e1fc05bab73304c3d12ce0c640cd7dfd Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Sun, 21 Jul 2019 02:59:01 +0100
Subject: [PATCH] create dvc.progress.Tqdm, migrate dvc.remote.base

---
 dvc/progress.py    | 41 +++++++++++++++++++++++++++++++++--------
 dvc/remote/base.py | 24 +++++-------------------
 2 files changed, 38 insertions(+), 27 deletions(-)

diff --git a/dvc/progress.py b/dvc/progress.py
index 81d4ba658..1f4d02965 100644
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -1,25 +1,50 @@
 """Manages progress bars for dvc repo."""
-
 from __future__ import print_function
-from __future__ import unicode_literals
 import logging
 from tqdm import tqdm
+from threading import Lock
 
-from dvc.utils.compat import str
 
-import sys
-import threading
+class Tqdm(tqdm):
+    """
+    maximum-compatibility tqdm-based progressbars
+    """
+    def __init__(
+            self,
+            iterable=None,
+            disable=logging.getLogger(
+                __name__).getEffectiveLevel() >= logging.CRITICAL,
+            ascii=None,  # TODO: True?
+            **kwargs):
+        """
+        kwargs  : anything accepted by `tqdm.tqdm()`
+        """
+        super(Tqdm, self).__init__(
+            iterable=iterable,
+            disable=disable,
+            ascii=ascii,
+            **kwargs)
+        # self.set_lock(Lock())
+
+    def update_desc(self, desc, n=1):
+        """
+        Calls `set_description(desc)` and `update(n)`
+        """
+        self.set_description(desc, refresh=False)
+        self.update(n)
 
 
 class Progress(tqdm):
     """
     Simple multi-target progress bar.
+    TODO: remove this class.
     """
     def __init__(self):
         super(Progress, self).__init__(
             total=0,
-            disable=logging.getLogger(__name__).getEffectiveLevel() >= logging.CRITICAL)
-        self.set_lock(threading.Lock())
+            disable=logging.getLogger(
+                __name__).getEffectiveLevel() >= logging.CRITICAL)
+        self.set_lock(Lock())
         self._targets = {}
         self.clearln()
 
@@ -77,7 +102,7 @@ class Progress(tqdm):
 
 
 class ProgressCallback(object):
-    # TODO: is this meant to be a thin wrapper for multi-progress, or just one bar?
+    # TODO: remove this thin wrapper
     def __init__(self, total):
         self.total = total
         self.current = 0
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index 9a387774a..224744d52 100644
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -19,7 +19,7 @@ from dvc.exceptions import (
     ConfirmRemoveError,
     DvcIgnoreInCollectedDirError,
 )
-from dvc.progress import progress, ProgressCallback
+from dvc.progress import Tqdm
 from dvc.utils import (
     LARGE_DIR_SIZE,
     tmp_fname,
@@ -159,7 +159,7 @@ class RemoteBASE(object):
                     "This is only done once."
                 )
                 logger.info(msg)
-                tasks = progress(tasks, total=len(file_infos))
+                tasks = Tqdm(tasks, total=len(file_infos), unit="md5")
 
         checksums = {
             file_infos[index]: task for index, task in enumerate(tasks)
@@ -444,9 +444,6 @@ class RemoteBASE(object):
 
         name = name or from_info.name
 
-        if not no_progress_bar:
-            progress.update_target(name, 0, None)
-
         try:
             self._upload(
                 from_info.fspath,
@@ -459,9 +456,6 @@ class RemoteBASE(object):
             logger.exception(msg.format(from_info, to_info))
             return 1  # 1 fail
 
-        if not no_progress_bar:
-            progress.finish_target(name)
-
         return 0
 
     def download(
@@ -490,11 +484,6 @@ class RemoteBASE(object):
 
         name = name or to_info.name
 
-        if not no_progress_bar:
-            # real progress is not always available,
-            # lets at least show start and finish
-            progress.update_target(name, 0, None)
-
         makedirs(to_info.parent, exist_ok=True, mode=dir_mode)
         tmp_file = tmp_fname(to_info)
 
@@ -509,9 +498,6 @@ class RemoteBASE(object):
 
         move(tmp_file, to_info, mode=file_mode)
 
-        if not no_progress_bar:
-            progress.finish_target(name)
-
         return 0
 
     def open(self, path_info, mode="r", encoding=None):
@@ -642,10 +628,10 @@ class RemoteBASE(object):
         Returns:
             A list with checksums that were found in the remote
         """
-        progress_callback = ProgressCallback(len(checksums))
+        pbar = Tqdm(total=len(checksums))
 
         def exists_with_progress(chunks):
-            return self.batch_exists(chunks, callback=progress_callback)
+            return self.batch_exists(chunks, callback=pbar.update_desc)
 
         if self.no_traverse and hasattr(self, "batch_exists"):
             with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor:
@@ -654,7 +640,7 @@ class RemoteBASE(object):
                 results = executor.map(exists_with_progress, chunks)
                 in_remote = itertools.chain.from_iterable(results)
                 ret = list(itertools.compress(checksums, in_remote))
-                progress_callback.finish("")
+                pbar.close()
                 return ret
 
         return list(set(checksums) & set(self.all()))