def download_import_files(self, task_id, data_dir):
import_struct = self.api.simple_request('GetImportStructure', sly.api_proto.ListFiles,
sly.api_proto.Id(id=task_id))
progress = sly.Progress('Downloading', len(import_struct.files), self.logger)
def close_fh(fh):
fpath = fh.file_path
if fh.close_and_check():
progress.iter_done_report()
else:
self.logger.warning('file was skipped while downloading', extra={'file_path': fpath})
file_handler = None
for chunk in self.api.get_stream_with_data('GetImportFiles',
sly.api_proto.ChunkFile,
sly.api_proto.ImportRequest(task_id=task_id,
files=import_struct.files)):
new_fpath = chunk.file.path
if new_fpath: # non-empty
if file_handler is not None:
close_fh(file_handler)
real_fpath = os.path.join(data_dir, new_fpath.lstrip('/'))
self.logger.trace('download import file', extra={'file_path': real_fpath})
file_handler = sly.ChunkedFileWriter(file_path=real_fpath)
file_handler.write(chunk.chunk)
close_fh(file_handler)
def download_import_files(self, task_id, data_dir):
import_struct = self.api.simple_request('GetImportStructure',
sly.api_proto.ListFiles,
sly.api_proto.Id(id=task_id))
progress = sly.Progress('Downloading', len(import_struct.files),
self.logger)
def maybe_close_fh(fh, pbar, downloaded_paths: set):
if fh is not None:
if fh.close_and_check():
pbar.iter_done_report()
downloaded_paths.add(fh.file_path)
else:
self.logger.warning('file was skipped while downloading',
extra={'file_path': fh.file_path})
files_to_download = list(import_struct.files)
for batch in batched(files_to_download):
# Store the file names that have been already downloaded from this batch
# to avoid rewriting them on transmission retries if connection issues arise.
downloaded_from_batch = set()
file_handler = None
for chunk in self.api.get_stream_with_data(
'GetImportFiles', sly.api_proto.ChunkFile,
sly.api_proto.ImportRequest(task_id=task_id, files=batch)):
new_fpath = chunk.file.path
if new_fpath: # non-empty
maybe_close_fh(file_handler, progress,
downloaded_from_batch)
real_fpath = os.path.join(data_dir, new_fpath.lstrip('/'))
if real_fpath in downloaded_from_batch:
file_handler = None
else:
self.logger.trace('download import file',
extra={'file_path': real_fpath})
file_handler = sly.ChunkedFileWriter(
file_path=real_fpath)
if file_handler is not None:
file_handler.write(chunk.chunk)
maybe_close_fh(file_handler, progress, downloaded_from_batch)