gdrive: add progress

Part of iterative#2865 See iterative#2865 (comment)
casperdcl · May 5, 2020 · b6d84f8 · b6d84f8
1 parent 274fa1c
commit b6d84f8
Showing 1 changed file with 28 additions and 8 deletions.
diff --git a/dvc/remote/gdrive.py b/dvc/remote/gdrive.py
@@ -396,14 +396,34 @@ def _gdrive_download_file(
         param = {"id": item_id}
         # it does not create a file on the remote
         gdrive_file = self._drive.CreateFile(param)
-        bar_format = (
-            "Downloading {desc:{ncols_desc}.{ncols_desc}}... "
-            + Tqdm.format_sizeof(int(gdrive_file["fileSize"]), "B", 1024)
-        )
-        with Tqdm(
-            bar_format=bar_format, desc=progress_desc, disable=no_progress_bar
-        ):
-            gdrive_file.GetContentFile(to_file)  # TODO: actually use pbar
+
+        import httplib2
+
+        OrigClass = httplib2.HTTPConnectionWithTimeout.response_class
+
+        class Custom(OrigClass):
+            def _readall_chunked(self):
+                assert self.chunked != "UNKNOWN"
+                value = []
+                with Tqdm(
+                    total=int(gdrive_file["fileSize"]),
+                    desc=progress_desc,
+                    disable=no_progress_bar,
+                    bytes=True,
+                ) as pbar:
+                    while True:
+                        chunk_left = self._get_chunk_left()
+                        if chunk_left is None:
+                            break
+                        chunk = self._safe_read(chunk_left)
+                        value.append(chunk)
+                        pbar.update(len(chunk))
+                        self.chunk_left = 0
+                return b"".join(value)
+
+        httplib2.HTTPConnectionWithTimeout.response_class = Custom
+        gdrive_file.GetContentFile(to_file)
+        httplib2.HTTPConnectionWithTimeout.response_class = OrigClass
 
     @_gdrive_retry
     def _gdrive_delete_file(self, item_id):