@@ -84,23 +84,9 @@ def init_drive(self):
84
84
)
85
85
)
86
86
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 )
88
88
self .cached_dirs , self .cached_ids = self .cache_root_dirs ()
89
89
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
-
104
90
def gdrive_upload_file (
105
91
self , args , no_progress_bar = True , from_file = "" , progress_name = ""
106
92
):
@@ -124,11 +110,8 @@ def gdrive_download_file(
124
110
from dvc .progress import Tqdm
125
111
126
112
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 )
132
115
133
116
def gdrive_list_item (self , query ):
134
117
file_list = self .drive .ListFile ({"q" : query , "maxResults" : 1000 })
@@ -192,12 +175,20 @@ def drive(self):
192
175
gdrive = GoogleDrive (gauth )
193
176
return gdrive
194
177
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
199
189
200
- def get_drive_item (self , name , parents_ids ):
190
+ @gdrive_retry
191
+ def get_remote_item (self , name , parents_ids ):
201
192
if not parents_ids :
202
193
return None
203
194
query = " or " .join (
@@ -206,14 +197,15 @@ def get_drive_item(self, name, parents_ids):
206
197
207
198
query += " and trashed=false and title='{}'" .format (name )
208
199
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 ()
210
202
return next (iter (item_list ), None )
211
203
212
- def resolve_remote_file (self , parents_ids , path_parts , create ):
204
+ def resolve_remote_item_from_path (self , parents_ids , path_parts , create ):
213
205
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 )
215
207
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 )
217
209
elif not item :
218
210
return None
219
211
parents_ids = [item ["id" ]]
@@ -230,7 +222,7 @@ def subtract_root_path(self, parts):
230
222
break
231
223
return parts , [self .root_id ]
232
224
233
- def get_path_id_from_cache (self , path_info ):
225
+ def get_remote_id_from_cache (self , path_info ):
234
226
files_ids = []
235
227
parts , parents_ids = self .subtract_root_path (path_info .path .split ("/" ))
236
228
if (
@@ -245,22 +237,22 @@ def get_path_id_from_cache(self, path_info):
245
237
246
238
return files_ids , parents_ids , parts
247
239
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 )
250
242
251
243
if not parts and files_ids :
252
244
return files_ids [0 ]
253
245
254
- file1 = self .resolve_remote_file (parents_ids , parts , create )
246
+ file1 = self .resolve_remote_item_from_path (parents_ids , parts , create )
255
247
return file1 ["id" ] if file1 else ""
256
248
257
249
def exists (self , path_info ):
258
- return self .get_path_id (path_info ) != ""
250
+ return self .get_remote_id (path_info ) != ""
259
251
260
252
def _upload (self , from_file , to_info , name , no_progress_bar ):
261
253
dirname = to_info .parent
262
254
if dirname :
263
- parent_id = self .get_path_id (dirname , True )
255
+ parent_id = self .get_remote_id (dirname , True )
264
256
else :
265
257
parent_id = to_info .bucket
266
258
@@ -274,33 +266,33 @@ def _upload(self, from_file, to_info, name, no_progress_bar):
274
266
)()
275
267
276
268
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 )
278
270
gdrive_retry (
279
271
lambda : self .gdrive_download_file (
280
272
file_id , to_file , name , no_progress_bar
281
273
)
282
274
)()
283
275
284
276
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 )
286
278
prefix = self .path_info .path
287
- for path in self .list_path (file_id ):
279
+ for path in self .list_children (file_id ):
288
280
yield posixpath .join (prefix , path )
289
281
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 ):
298
283
for file1 in self .gdrive_list_item (
299
284
"'{}' in parents and trashed=false" .format (parent_id )
300
285
):
301
- for path in self .list_file_path (file1 ):
286
+ for path in self .list_remote_item (file1 ):
302
287
yield path
303
288
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
+
304
296
def all (self ):
305
297
if not hasattr (self , "cached_ids" ) or not self .cached_ids :
306
298
return
0 commit comments