Skip to content

Commit

Permalink
clean up gs progress using Tqdm.wrapattr
Browse files Browse the repository at this point in the history
- fixes #2829
  • Loading branch information
casperdcl committed Dec 1, 2019
1 parent f09517a commit ecdc608
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 33 deletions.
50 changes: 18 additions & 32 deletions dvc/remote/gs.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,23 +58,15 @@ def _upload_to_bucket(
no_progress_bar=True,
):
blob = bucket.blob(to_info.path, chunk_size=chunk_size)
with Tqdm(
desc=name or to_info.path,
total=os.path.getsize(from_file),
bytes=True,
disable=no_progress_bar,
) as pbar:
with io.open(from_file, mode="rb") as fobj:
raw_read = fobj.read

def read(size=chunk_size):
res = raw_read(size)
if res:
pbar.update(len(res))
return res

fobj.read = read
blob.upload_from_file(fobj)
with io.open(from_file, mode="rb") as fobj:
with Tqdm.wrapattr(
fobj,
"read",
desc=name or to_info.path,
total=os.path.getsize(from_file),
disable=no_progress_bar,
) as wrapped:
blob.upload_from_file(wrapped)


class RemoteGS(RemoteBASE):
Expand Down Expand Up @@ -162,21 +154,15 @@ def _upload(self, from_file, to_info, name=None, no_progress_bar=True):
def _download(self, from_info, to_file, name=None, no_progress_bar=True):
bucket = self.gs.bucket(from_info.bucket)
blob = bucket.get_blob(from_info.path)
with Tqdm(
desc=name or from_info.path,
total=blob.size,
bytes=True,
disable=no_progress_bar,
) as pbar:
with io.open(to_file, mode="wb") as fobj:
raw_write = fobj.write

def write(byte_string):
raw_write(byte_string)
pbar.update(len(byte_string))

fobj.write = write
blob.download_to_file(fobj)
with io.open(to_file, mode="wb") as fobj:
with Tqdm.wrapattr(
fobj,
"write",
desc=name or from_info.path,
total=blob.size,
disable=no_progress_bar,
) as wrapped:
blob.download_to_file(wrapped)

def _generate_download_url(self, path_info, expires=3600):
expiration = timedelta(seconds=int(expires))
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ def run(self):
"funcy>=1.14",
"pathspec>=0.6.0",
"shortuuid>=0.5.0",
"tqdm>=4.38.0,<5",
"tqdm>=4.40.0,<5",
"packaging>=19.0",
"win-unicode-console>=0.5; sys_platform == 'win32'",
"pywin32>=225; sys_platform == 'win32'",
Expand Down

0 comments on commit ecdc608

Please sign in to comment.