Skip to content

Commit ecdc608

Browse files
committed
clean up gs progress using Tqdm.wrapattr
- fixes #2829
1 parent f09517a commit ecdc608

File tree

2 files changed

+19
-33
lines changed

2 files changed

+19
-33
lines changed

dvc/remote/gs.py

+18-32
Original file line numberDiff line numberDiff line change
@@ -58,23 +58,15 @@ def _upload_to_bucket(
5858
no_progress_bar=True,
5959
):
6060
blob = bucket.blob(to_info.path, chunk_size=chunk_size)
61-
with Tqdm(
62-
desc=name or to_info.path,
63-
total=os.path.getsize(from_file),
64-
bytes=True,
65-
disable=no_progress_bar,
66-
) as pbar:
67-
with io.open(from_file, mode="rb") as fobj:
68-
raw_read = fobj.read
69-
70-
def read(size=chunk_size):
71-
res = raw_read(size)
72-
if res:
73-
pbar.update(len(res))
74-
return res
75-
76-
fobj.read = read
77-
blob.upload_from_file(fobj)
61+
with io.open(from_file, mode="rb") as fobj:
62+
with Tqdm.wrapattr(
63+
fobj,
64+
"read",
65+
desc=name or to_info.path,
66+
total=os.path.getsize(from_file),
67+
disable=no_progress_bar,
68+
) as wrapped:
69+
blob.upload_from_file(wrapped)
7870

7971

8072
class RemoteGS(RemoteBASE):
@@ -162,21 +154,15 @@ def _upload(self, from_file, to_info, name=None, no_progress_bar=True):
162154
def _download(self, from_info, to_file, name=None, no_progress_bar=True):
163155
bucket = self.gs.bucket(from_info.bucket)
164156
blob = bucket.get_blob(from_info.path)
165-
with Tqdm(
166-
desc=name or from_info.path,
167-
total=blob.size,
168-
bytes=True,
169-
disable=no_progress_bar,
170-
) as pbar:
171-
with io.open(to_file, mode="wb") as fobj:
172-
raw_write = fobj.write
173-
174-
def write(byte_string):
175-
raw_write(byte_string)
176-
pbar.update(len(byte_string))
177-
178-
fobj.write = write
179-
blob.download_to_file(fobj)
157+
with io.open(to_file, mode="wb") as fobj:
158+
with Tqdm.wrapattr(
159+
fobj,
160+
"write",
161+
desc=name or from_info.path,
162+
total=blob.size,
163+
disable=no_progress_bar,
164+
) as wrapped:
165+
blob.download_to_file(wrapped)
180166

181167
def _generate_download_url(self, path_info, expires=3600):
182168
expiration = timedelta(seconds=int(expires))

setup.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ def run(self):
7777
"funcy>=1.14",
7878
"pathspec>=0.6.0",
7979
"shortuuid>=0.5.0",
80-
"tqdm>=4.38.0,<5",
80+
"tqdm>=4.40.0,<5",
8181
"packaging>=19.0",
8282
"win-unicode-console>=0.5; sys_platform == 'win32'",
8383
"pywin32>=225; sys_platform == 'win32'",

0 commit comments

Comments
 (0)