Merge pull request #95 from efiop/master

progress: add progress bar to 'import'

dvc/command/import_file.py

@@ -12,6 +12,7 @@ from dvc.runtime import Runtime
 from dvc.state_file import StateFile
 from dvc.system import System
 from dvc.command.data_sync import POOL_SIZE
+from dvc.progress import Progress
 
 class ImportFileError(DvcException):
     def __init__(self, msg):
@@ -140,27 +141,41 @@ class CmdImportFile(CmdBase):
         """
         r = requests.get(from_url, stream=True)
 
+        name = os.path.basename(from_url)
         chunk_size = 1024 * 100
         downloaded = 0
         last_reported = 0
         report_bucket = 100*1024*10
+        total_length = r.headers.get('content-length')
 
         with open(to_file, 'wb') as f:
             for chunk in r.iter_content(chunk_size=chunk_size):
-                if chunk:  # filter out keep-alive new chunks
-                    downloaded += chunk_size
-                    last_reported += chunk_size
-                    if last_reported >= report_bucket:
-                        last_reported = 0
-                        Logger.debug('Downloaded {}'.format(sizeof_fmt(downloaded)))
-                    f.write(chunk)
+                if not chunk:  # filter out keep-alive new chunks
+                    continue
+
+                downloaded += chunk_size
+
+                last_reported += chunk_size
+                if last_reported >= report_bucket:
+                    last_reported = 0
+                    Logger.debug('Downloaded {}'.format(sizeof_fmt(downloaded)))
+
+                # update progress bar
+                self.progress.update_target(name, downloaded, total_length)
+
+                f.write(chunk)
+
+        # tell progress bar that this target is finished downloading
+        self.progress.finish_target(name)
 
     def download_targets(self, targets):
         """
         Download targets in a number of threads.
         """
+        self.progress = Progress(len(targets))
         p = ThreadPool(processes=POOL_SIZE)
         p.map(self.download_target, targets)
+        self.progress.finish()
 
     def create_state_files(self, targets, lock):
         """

dvc/progress.py

@@ -0,0 +1,75 @@
+from __future__ import print_function
+import sys
+import threading 
+
+class Progress(object):
+    """
+    Simple multi-target progress bar.
+    """
+    def __init__(self, total):
+        self._n_total = total
+        self._n_finished = 0
+        self._lock = threading.Lock()
+
+    def _clearln(self):
+        print('\r\x1b[K', end='')
+
+    def _writeln(self, line):
+        self._clearln()
+        print(line, end='')
+        sys.stdout.flush()
+
+    def update_target(self, name, current, total):
+        # Just go away if it is locked. Will update next time
+        if not self._lock.acquire(False):
+            return
+
+        bar = self._bar(name, current, total)
+
+        if sys.stdout.isatty():
+            self._writeln(bar)
+
+        self._lock.release()
+
+    def finish_target(self, name):
+        # We have to write a msg about finished target
+        with self._lock:
+            bar = self._bar(name, 100, 100)
+
+            if sys.stdout.isatty():
+                self._clearln()
+
+            print(bar)
+
+            self._n_finished += 1
+
+    def finish(self):
+        if sys.stdout.isatty():
+            print()
+
+    def _bar(self, target_name, current, total):
+        """
+        Make a progress bar out of info, which looks like:
+        (1/2): [########################################] 100% master.zip
+        """
+        total_len = 80
+        bar_len = 40
+
+        if total == None:
+            progress = 0
+            percent = "?% "
+        else:
+            total = int(total)
+            progress = int((100 * current)/total) if current <= total else 100
+            percent = str(progress) + "% "
+
+        num = "({}/{}): ".format(self._n_finished + 1, self._n_total)
+
+        n_sh = int((progress * bar_len)/100)
+        n_sp = bar_len - n_sh
+        bar = "[" + '#'*n_sh + ' '*n_sp + "] "
+
+        name_len = total_len - len(num + bar + percent)
+        name = target_name[:name_len] if len(target_name) > name_len else target_name
+
+        return num + bar + percent + name