From 62cbeed8b38bc52bf72a3b42b77c67e6fae1a5a7 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Sun, 28 Jul 2019 14:53:02 +0100 Subject: [PATCH] tqdm remote.http --- dvc/remote/http.py | 49 +++++++++++++++------------------------------- 1 file changed, 16 insertions(+), 33 deletions(-) diff --git a/dvc/remote/http.py b/dvc/remote/http.py index f39c5a3e3..911e9754c 100644 --- a/dvc/remote/http.py +++ b/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