diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 6f389ec4f..b1d5c960f 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -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): @@ -162,21 +154,15 @@ class RemoteGS(RemoteBASE): 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)) diff --git a/setup.py b/setup.py index b9b16f014..b23d269db 100644 --- a/setup.py +++ b/setup.py @@ -77,7 +77,7 @@ install_requires = [ "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'",