tqdm remote.http

dvc/remote/http.py

@@ -1,35 +1,20 @@
 from __future__ import unicode_literals
 
 from dvc.scheme import Schemes
-
 from dvc.utils.compat import open
 
 import threading
 import requests
 import logging
 
-from dvc.progress import progress
+from dvc.progress import Tqdm
 from dvc.exceptions import DvcException
 from dvc.config import Config
 from dvc.remote.base import RemoteBASE
 
-
 logger = logging.getLogger(__name__)
 
 
-class ProgressBarCallback(object):
-    def __init__(self, name, total):
-        self.name = name
-        self.total = total
-        self.current = 0
-        self.lock = threading.Lock()
-
-    def __call__(self, byts):
-        with self.lock:
-            self.current += byts
-            progress.update_target(self.name, self.current, self.total)
-
-
 class RemoteHTTP(RemoteBASE):
     scheme = Schemes.HTTP
     REQUEST_TIMEOUT = 10
@@ -43,24 +28,21 @@ class RemoteHTTP(RemoteBASE):
         self.path_info = self.path_cls(url) if url else None
 
     def _download(self, from_info, to_file, name=None, no_progress_bar=False):
-        callback = None
-        if not no_progress_bar:
-            total = self._content_length(from_info.url)
-            if total:
-                callback = ProgressBarCallback(name, total)
-
         request = self._request("GET", from_info.url, stream=True)
+        total = self._content_length(from_info.url)
+        leave = False
+        #TODO: persistent progress only for "large" files?
+        #if total:
+        #    leave = total > self.CHUNK_SIZE * 100
 
-        with open(to_file, "wb") as fd:
-            transferred_bytes = 0
-
-            for chunk in request.iter_content(chunk_size=self.CHUNK_SIZE):
-                fd.write(chunk)
-                fd.flush()
-                transferred_bytes += len(chunk)
-
-                if callback:
-                    callback(transferred_bytes)
+        with Tqdm(total=total, leave=leave,
+                  unit='B', unit_scale=True, unit_divisor=1024, miniters=1,
+                  disable=no_progress_bar) as pbar:
+            with open(to_file, "wb") as fd:
+                for chunk in request.iter_content(chunk_size=self.CHUNK_SIZE):
+                    fd.write(chunk)
+                    fd.flush()
+                    pbar.update(len(chunk))
 
     def exists(self, path_info):
         return bool(self._request("HEAD", path_info.url))
@@ -75,7 +57,8 @@ class RemoteHTTP(RemoteBASE):
         return results
 
     def _content_length(self, url):
-        return self._request("HEAD", url).headers.get("Content-Length")
+        res = self._request("HEAD", url).headers.get("Content-Length")
+        return int(res) if res else None
 
     def get_file_checksum(self, path_info):
         url = path_info.url