Skip to content

Commit b6d84f8

Browse files
committed
gdrive: add progress
Part of iterative#2865 See iterative#2865 (comment)
1 parent 274fa1c commit b6d84f8

File tree

1 file changed

+28
-8
lines changed

1 file changed

+28
-8
lines changed

dvc/remote/gdrive.py

Lines changed: 28 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -396,14 +396,34 @@ def _gdrive_download_file(
396396
param = {"id": item_id}
397397
# it does not create a file on the remote
398398
gdrive_file = self._drive.CreateFile(param)
399-
bar_format = (
400-
"Downloading {desc:{ncols_desc}.{ncols_desc}}... "
401-
+ Tqdm.format_sizeof(int(gdrive_file["fileSize"]), "B", 1024)
402-
)
403-
with Tqdm(
404-
bar_format=bar_format, desc=progress_desc, disable=no_progress_bar
405-
):
406-
gdrive_file.GetContentFile(to_file) # TODO: actually use pbar
399+
400+
import httplib2
401+
402+
OrigClass = httplib2.HTTPConnectionWithTimeout.response_class
403+
404+
class Custom(OrigClass):
405+
def _readall_chunked(self):
406+
assert self.chunked != "UNKNOWN"
407+
value = []
408+
with Tqdm(
409+
total=int(gdrive_file["fileSize"]),
410+
desc=progress_desc,
411+
disable=no_progress_bar,
412+
bytes=True,
413+
) as pbar:
414+
while True:
415+
chunk_left = self._get_chunk_left()
416+
if chunk_left is None:
417+
break
418+
chunk = self._safe_read(chunk_left)
419+
value.append(chunk)
420+
pbar.update(len(chunk))
421+
self.chunk_left = 0
422+
return b"".join(value)
423+
424+
httplib2.HTTPConnectionWithTimeout.response_class = Custom
425+
gdrive_file.GetContentFile(to_file)
426+
httplib2.HTTPConnectionWithTimeout.response_class = OrigClass
407427

408428
@_gdrive_retry
409429
def _gdrive_delete_file(self, item_id):

0 commit comments

Comments
 (0)