diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py index a97911364..59f2f8a2c 100644 --- a/dvc/command/import_file.py +++ b/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): """ diff --git a/dvc/progress.py b/dvc/progress.py new file mode 100644 index 000000000..31c45b682 --- /dev/null +++ b/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