Skip to content

Commit b883bb9

Browse files
committed
Use gdrive_retry as decorator;Use Tqdm with disable param; Rename methods.
1 parent 5294cff commit b883bb9

File tree

1 file changed

+39
-47
lines changed

1 file changed

+39
-47
lines changed

dvc/remote/gdrive/__init__.py

Lines changed: 39 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -84,23 +84,9 @@ def init_drive(self):
8484
)
8585
)
8686

87-
self.root_id = self.get_path_id(self.path_info, create=True)
87+
self.root_id = self.get_remote_id(self.path_info, create=True)
8888
self.cached_dirs, self.cached_ids = self.cache_root_dirs()
8989

90-
def gdrive_list_file(self, query):
91-
return self.drive.ListFile({"q": query, "maxResults": 1}).GetList()
92-
93-
def gdrive_create_folder(self, title, parent_id):
94-
item = self.drive.CreateFile(
95-
{
96-
"title": title,
97-
"parents": [{"id": parent_id}],
98-
"mimeType": FOLDER_MIME_TYPE,
99-
}
100-
)
101-
item.Upload()
102-
return item
103-
10490
def gdrive_upload_file(
10591
self, args, no_progress_bar=True, from_file="", progress_name=""
10692
):
@@ -124,11 +110,8 @@ def gdrive_download_file(
124110
from dvc.progress import Tqdm
125111

126112
gdrive_file = self.drive.CreateFile({"id": file_id})
127-
if not no_progress_bar:
128-
tqdm = Tqdm(desc=progress_name, total=int(gdrive_file["fileSize"]))
129-
gdrive_file.GetContentFile(to_file)
130-
if not no_progress_bar:
131-
tqdm.close()
113+
with Tqdm(desc=progress_name, total=int(gdrive_file["fileSize"]), disable=no_progress_bar):
114+
gdrive_file.GetContentFile(to_file)
132115

133116
def gdrive_list_item(self, query):
134117
file_list = self.drive.ListFile({"q": query, "maxResults": 1000})
@@ -192,12 +175,20 @@ def drive(self):
192175
gdrive = GoogleDrive(gauth)
193176
return gdrive
194177

195-
def create_drive_item(self, parent_id, title):
196-
return gdrive_retry(
197-
lambda: self.gdrive_create_folder(title, parent_id)
198-
)()
178+
@gdrive_retry
179+
def create_remote_dir(self, parent_id, title):
180+
item = self.drive.CreateFile(
181+
{
182+
"title": title,
183+
"parents": [{"id": parent_id}],
184+
"mimeType": FOLDER_MIME_TYPE,
185+
}
186+
)
187+
item.Upload()
188+
return item
199189

200-
def get_drive_item(self, name, parents_ids):
190+
@gdrive_retry
191+
def get_remote_item(self, name, parents_ids):
201192
if not parents_ids:
202193
return None
203194
query = " or ".join(
@@ -206,14 +197,15 @@ def get_drive_item(self, name, parents_ids):
206197

207198
query += " and trashed=false and title='{}'".format(name)
208199

209-
item_list = gdrive_retry(lambda: self.gdrive_list_file(query))()
200+
# Limit found remote items count to 1 in response
201+
item_list = self.drive.ListFile({"q": query, "maxResults": 1}).GetList()
210202
return next(iter(item_list), None)
211203

212-
def resolve_remote_file(self, parents_ids, path_parts, create):
204+
def resolve_remote_item_from_path(self, parents_ids, path_parts, create):
213205
for path_part in path_parts:
214-
item = self.get_drive_item(path_part, parents_ids)
206+
item = self.get_remote_item(path_part, parents_ids)
215207
if not item and create:
216-
item = self.create_drive_item(parents_ids[0], path_part)
208+
item = self.create_remote_dir(parents_ids[0], path_part)
217209
elif not item:
218210
return None
219211
parents_ids = [item["id"]]
@@ -230,7 +222,7 @@ def subtract_root_path(self, parts):
230222
break
231223
return parts, [self.root_id]
232224

233-
def get_path_id_from_cache(self, path_info):
225+
def get_remote_id_from_cache(self, path_info):
234226
files_ids = []
235227
parts, parents_ids = self.subtract_root_path(path_info.path.split("/"))
236228
if (
@@ -245,22 +237,22 @@ def get_path_id_from_cache(self, path_info):
245237

246238
return files_ids, parents_ids, parts
247239

248-
def get_path_id(self, path_info, create=False):
249-
files_ids, parents_ids, parts = self.get_path_id_from_cache(path_info)
240+
def get_remote_id(self, path_info, create=False):
241+
files_ids, parents_ids, parts = self.get_remote_id_from_cache(path_info)
250242

251243
if not parts and files_ids:
252244
return files_ids[0]
253245

254-
file1 = self.resolve_remote_file(parents_ids, parts, create)
246+
file1 = self.resolve_remote_item_from_path(parents_ids, parts, create)
255247
return file1["id"] if file1 else ""
256248

257249
def exists(self, path_info):
258-
return self.get_path_id(path_info) != ""
250+
return self.get_remote_id(path_info) != ""
259251

260252
def _upload(self, from_file, to_info, name, no_progress_bar):
261253
dirname = to_info.parent
262254
if dirname:
263-
parent_id = self.get_path_id(dirname, True)
255+
parent_id = self.get_remote_id(dirname, True)
264256
else:
265257
parent_id = to_info.bucket
266258

@@ -274,33 +266,33 @@ def _upload(self, from_file, to_info, name, no_progress_bar):
274266
)()
275267

276268
def _download(self, from_info, to_file, name, no_progress_bar):
277-
file_id = self.get_path_id(from_info)
269+
file_id = self.get_remote_id(from_info)
278270
gdrive_retry(
279271
lambda: self.gdrive_download_file(
280272
file_id, to_file, name, no_progress_bar
281273
)
282274
)()
283275

284276
def list_cache_paths(self):
285-
file_id = self.get_path_id(self.path_info)
277+
file_id = self.get_remote_id(self.path_info)
286278
prefix = self.path_info.path
287-
for path in self.list_path(file_id):
279+
for path in self.list_children(file_id):
288280
yield posixpath.join(prefix, path)
289281

290-
def list_file_path(self, drive_file):
291-
if drive_file["mimeType"] == FOLDER_MIME_TYPE:
292-
for i in self.list_path(drive_file["id"]):
293-
yield posixpath.join(drive_file["title"], i)
294-
else:
295-
yield drive_file["title"]
296-
297-
def list_path(self, parent_id):
282+
def list_children(self, parent_id):
298283
for file1 in self.gdrive_list_item(
299284
"'{}' in parents and trashed=false".format(parent_id)
300285
):
301-
for path in self.list_file_path(file1):
286+
for path in self.list_remote_item(file1):
302287
yield path
303288

289+
def list_remote_item(self, drive_file):
290+
if drive_file["mimeType"] == FOLDER_MIME_TYPE:
291+
for i in self.list_children(drive_file["id"]):
292+
yield posixpath.join(drive_file["title"], i)
293+
else:
294+
yield drive_file["title"]
295+
304296
def all(self):
305297
if not hasattr(self, "cached_ids") or not self.cached_ids:
306298
return

0 commit comments

Comments
 (0)