From f61e7d86a99b09a7bf7ef2d84d652b57b538e286 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Sun, 28 Jul 2019 15:19:17 +0100 Subject: [PATCH] tqdm remote.azure --- dvc/progress.py | 5 +++++ dvc/remote/azure.py | 30 +++++++++++++----------------- dvc/remote/http.py | 2 +- 3 files changed, 19 insertions(+), 18 deletions(-) diff --git a/dvc/progress.py b/dvc/progress.py index 051ad855b..4d690b2a5 100644 --- a/dvc/progress.py +++ b/dvc/progress.py @@ -40,6 +40,11 @@ class Tqdm(tqdm): self.set_description(desc, refresh=False) self.update(n) + def update_to(self, current, total=None): + if total: + self.total = total + self.update(current - self.n) + @classmethod def truncate(cls, s, max_len=10, end=True, fill="..."): """ diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 0b8330ae3..af756aae6 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -15,7 +15,7 @@ except ImportError: BlockBlobService = None from dvc.utils.compat import urlparse -from dvc.progress import progress +from dvc.progress import Tqdm from dvc.config import Config from dvc.remote.base import RemoteBASE from dvc.path_info import CloudURLInfo @@ -25,14 +25,6 @@ from dvc.utils.http import open_url logger = logging.getLogger(__name__) -class Callback(object): - def __init__(self, name): - self.name = name - - def __call__(self, current, total): - progress.update_target(self.name, current, total) - - class RemoteAZURE(RemoteBASE): scheme = Schemes.AZURE path_cls = CloudURLInfo @@ -124,18 +116,22 @@ class RemoteAZURE(RemoteBASE): def _upload( self, from_file, to_info, name=None, no_progress_bar=False, **_kwargs ): - cb = None if no_progress_bar else Callback(name) - self.blob_service.create_blob_from_path( - to_info.bucket, to_info.path, from_file, progress_callback=cb - ) + desc = Tqdm.truncate(name, 10) if name else None + with Tqdm(desc=desc, disable=no_progress_bar) as pbar: + self.blob_service.create_blob_from_path( + to_info.bucket, to_info.path, from_file, + progress_callback=pbar.update_to + ) def _download( self, from_info, to_file, name=None, no_progress_bar=False, **_kwargs ): - cb = None if no_progress_bar else Callback(name) - self.blob_service.get_blob_to_path( - from_info.bucket, from_info.path, to_file, progress_callback=cb - ) + desc = Tqdm.truncate(name, 10) if name else None + with Tqdm(desc=desc, disable=no_progress_bar) as pbar: + self.blob_service.get_blob_to_path( + from_info.bucket, from_info.path, to_file, + progress_callback=pbar.update_to + ) def open(self, path_info, mode="r", encoding=None): get_url = lambda: self._generate_download_url(path_info) # noqa: E731 diff --git a/dvc/remote/http.py b/dvc/remote/http.py index bd7bfdfd6..d125f5273 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -35,7 +35,7 @@ class RemoteHTTP(RemoteBASE): # leave = total > self.CHUNK_SIZE * 100 with Tqdm(total=total, leave=leave, bytes=True, - desc=Tqdm.truncate(to_file, 10), + desc=Tqdm.truncate(to_file), disable=no_progress_bar) as pbar: with open(to_file, "wb") as fd: for chunk in request.iter_content(chunk_size=self.CHUNK_SIZE):