From af6c4291b95e777bea10772fc7d0ca7385bad4c1 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Fri, 27 Mar 2020 20:17:34 +0000 Subject: [PATCH] progress: wrappers (#3522) * progress: add wrap_fn Closes #3521 Related https://github.com/tqdm/tqdm/issues/583 * progress: use wrappers in more places --- dvc/progress.py | 15 +++++++++++++ dvc/remote/base.py | 49 ++++++++++++++++++++++-------------------- dvc/remote/http.py | 53 ++++++++++++++++++++++++---------------------- dvc/utils/fs.py | 16 ++++++++------ 4 files changed, 79 insertions(+), 54 deletions(-) diff --git a/dvc/progress.py b/dvc/progress.py index a8ab61951..dd7341398 100644 --- a/dvc/progress.py +++ b/dvc/progress.py @@ -114,6 +114,21 @@ class Tqdm(tqdm): self.total = total # pylint: disable=W0613,W0201 self.update(current - self.n) + def wrap_fn(self, fn, callback=None): + """ + Returns a wrapped `fn` which calls `callback()` on each call. + `callback` is `self.update` by default. + """ + if callback is None: + callback = self.update + + def wrapped(*args, **kwargs): + res = fn(*args, **kwargs) + callback() + return res + + return wrapped + def close(self): if self.desc_persist is not None: self.set_description_str(self.desc_persist, refresh=False) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 8b36a5100..644e64c1e 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -176,14 +176,16 @@ class RemoteBASE(object): def _calculate_checksums(self, file_infos): file_infos = list(file_infos) - with ThreadPoolExecutor(max_workers=self.checksum_jobs) as executor: - tasks = executor.map(self.get_file_checksum, file_infos) - with Tqdm( - tasks, - total=len(file_infos), - unit="md5", - desc="Computing file/dir hashes (only done once)", - ) as tasks: + with Tqdm( + total=len(file_infos), + unit="md5", + desc="Computing file/dir hashes (only done once)", + ) as pbar: + worker = pbar.wrap_fn(self.get_file_checksum) + with ThreadPoolExecutor( + max_workers=self.checksum_jobs + ) as executor: + tasks = executor.map(worker, file_infos) checksums = dict(zip(file_infos, tasks)) return checksums @@ -609,22 +611,23 @@ class RemoteBASE(object): to_info / info.relative_to(from_info) for info in from_infos ) - with ThreadPoolExecutor(max_workers=self.JOBS) as executor: - download_files = partial( - self._download_file, - name=name, - no_progress_bar=True, - file_mode=file_mode, - dir_mode=dir_mode, + with Tqdm( + total=len(from_infos), + desc="Downloading directory", + unit="Files", + disable=no_progress_bar, + ) as pbar: + download_files = pbar.wrap_fn( + partial( + self._download_file, + name=name, + no_progress_bar=True, + file_mode=file_mode, + dir_mode=dir_mode, + ) ) - futures = executor.map(download_files, from_infos, to_infos) - with Tqdm( - futures, - total=len(from_infos), - desc="Downloading directory", - unit="Files", - disable=no_progress_bar, - ) as futures: + with ThreadPoolExecutor(max_workers=self.JOBS) as executor: + futures = executor.map(download_files, from_infos, to_infos) return sum(futures) def _download_file( diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 6ea016ca6..a40d9278d 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -55,39 +55,42 @@ class RemoteHTTP(RemoteBASE): response = self._request("GET", from_info.url, stream=True) if response.status_code != 200: raise HTTPError(response.status_code, response.reason) - with Tqdm( - total=None if no_progress_bar else self._content_length(response), - leave=False, - bytes=True, - desc=from_info.url if name is None else name, - disable=no_progress_bar, - ) as pbar: - with open(to_file, "wb") as fd: + with open(to_file, "wb") as fd: + with Tqdm.wrapattr( + fd, + "write", + total=None + if no_progress_bar + else self._content_length(response), + leave=False, + desc=from_info.url if name is None else name, + disable=no_progress_bar, + ) as fd_wrapped: for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE): - fd.write(chunk) - pbar.update(len(chunk)) + fd_wrapped.write(chunk) def _upload(self, from_file, to_info, name=None, no_progress_bar=False): - with Tqdm( - total=None if no_progress_bar else os.path.getsize(from_file), - leave=False, - bytes=True, - desc=to_info.url if name is None else name, - disable=no_progress_bar, - ) as pbar: - - def chunks(): - with open(from_file, "rb") as fd: + def chunks(): + with open(from_file, "rb") as fd: + with Tqdm.wrapattr( + fd, + "read", + total=None + if no_progress_bar + else os.path.getsize(from_file), + leave=False, + desc=to_info.url if name is None else name, + disable=no_progress_bar, + ) as fd_wrapped: while True: - chunk = fd.read(self.CHUNK_SIZE) + chunk = fd_wrapped.read(self.CHUNK_SIZE) if not chunk: break - pbar.update(len(chunk)) yield chunk - response = self._request("POST", to_info.url, data=chunks()) - if response.status_code not in (200, 201): - raise HTTPError(response.status_code, response.reason) + response = self._request("POST", to_info.url, data=chunks()) + if response.status_code not in (200, 201): + raise HTTPError(response.status_code, response.reason) def exists(self, path_info): return bool(self._request("HEAD", path_info.url)) diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index f4b9972d3..8bbc097c0 100644 --- a/dvc/utils/fs.py +++ b/dvc/utils/fs.py @@ -189,16 +189,20 @@ def copyfile(src, dest, no_progress_bar=False, name=None): try: System.reflink(src, dest) except DvcException: - with Tqdm( - desc=name, disable=no_progress_bar, total=total, bytes=True - ) as pbar: - with open(src, "rb") as fsrc, open(dest, "wb+") as fdest: + with open(src, "rb") as fsrc, open(dest, "wb+") as fdest: + with Tqdm.wrapattr( + fdest, + "write", + desc=name, + disable=no_progress_bar, + total=total, + bytes=True, + ) as fdest_wrapped: while True: buf = fsrc.read(LOCAL_CHUNK_SIZE) if not buf: break - fdest.write(buf) - pbar.update(len(buf)) + fdest_wrapped.write(buf) def walk_files(directory):