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

@@ -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)

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(

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))

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):