Skip to content

Commit b312895

Browse files
casperdclefiop
andauthored
progress: add main bars (#3594)
* progress: add main push/pull bar Fixes #3452 Related #1840 Related #3565 * adjust desc * fix merge Co-authored-by: Ruslan Kuprieiev <[email protected]> Co-authored-by: Ruslan Kuprieiev <[email protected]>
1 parent e2981d2 commit b312895

File tree

1 file changed

+38
-33
lines changed

1 file changed

+38
-33
lines changed

dvc/remote/local.py

Lines changed: 38 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -448,9 +448,11 @@ def _process(
448448
file_mode=self._file_mode,
449449
)
450450
status = STATUS_DELETED
451+
desc = "Downloading"
451452
else:
452453
func = remote.upload
453454
status = STATUS_NEW
455+
desc = "Uploading"
454456

455457
if jobs is None:
456458
jobs = remote.JOBS
@@ -466,42 +468,45 @@ def _process(
466468
dir_plans = self._get_plans(download, remote, dir_status, status)
467469
file_plans = self._get_plans(download, remote, file_status, status)
468470

469-
if len(dir_plans[0]) + len(file_plans[0]) == 0:
471+
total = len(dir_plans[0]) + len(file_plans[0])
472+
if total == 0:
470473
return 0
471474

472-
with ThreadPoolExecutor(max_workers=jobs) as executor:
473-
if download:
474-
fails = sum(executor.map(func, *dir_plans))
475-
fails += sum(executor.map(func, *file_plans))
476-
else:
477-
# for uploads, push files first, and any .dir files last
478-
479-
file_futures = {}
480-
for from_info, to_info, name in zip(*file_plans):
481-
file_futures[to_info] = executor.submit(
482-
func, from_info, to_info, name
483-
)
484-
dir_futures = {}
485-
for from_info, to_info, name in zip(*dir_plans):
486-
wait_futures = {
487-
future
488-
for file_path, future in file_futures.items()
489-
if file_path in dir_paths[to_info]
490-
}
491-
dir_futures[to_info] = executor.submit(
492-
self._dir_upload,
493-
func,
494-
wait_futures,
495-
from_info,
496-
to_info,
497-
name,
498-
)
499-
fails = sum(
500-
future.result()
501-
for future in concat(
502-
file_futures.values(), dir_futures.values()
475+
with Tqdm(total=total, unit="file", desc=desc) as pbar:
476+
func = pbar.wrap_fn(func)
477+
with ThreadPoolExecutor(max_workers=jobs) as executor:
478+
if download:
479+
fails = sum(executor.map(func, *dir_plans))
480+
fails += sum(executor.map(func, *file_plans))
481+
else:
482+
# for uploads, push files first, and any .dir files last
483+
484+
file_futures = {}
485+
for from_info, to_info, name in zip(*file_plans):
486+
file_futures[to_info] = executor.submit(
487+
func, from_info, to_info, name
488+
)
489+
dir_futures = {}
490+
for from_info, to_info, name in zip(*dir_plans):
491+
wait_futures = {
492+
future
493+
for file_path, future in file_futures.items()
494+
if file_path in dir_paths[to_info]
495+
}
496+
dir_futures[to_info] = executor.submit(
497+
self._dir_upload,
498+
func,
499+
wait_futures,
500+
from_info,
501+
to_info,
502+
name,
503+
)
504+
fails = sum(
505+
future.result()
506+
for future in concat(
507+
file_futures.values(), dir_futures.values()
508+
)
503509
)
504-
)
505510

506511
if fails:
507512
if download:

0 commit comments

Comments
 (0)