From cc281551eb3e10451c2b7f08d083a9c8dcf0b939 Mon Sep 17 00:00:00 2001 From: Casper da Costa-Luis Date: Sun, 28 Jul 2019 14:50:39 +0100 Subject: [PATCH] tqdm repo.checkout --- dvc/repo/checkout.py | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py index 2d743e41d..826a24a81 100644 --- a/dvc/repo/checkout.py +++ b/dvc/repo/checkout.py @@ -3,7 +3,7 @@ from __future__ import unicode_literals import logging from dvc.exceptions import CheckoutErrorSuggestGit -from dvc.progress import ProgressCallback +from dvc.progress import Tqdm logger = logging.getLogger(__name__) @@ -23,13 +23,6 @@ def get_all_files_numbers(stages): return sum(stage.get_all_files_number() for stage in stages) -def get_progress_callback(stages): - total_files_num = get_all_files_numbers(stages) - if total_files_num == 0: - return None - return ProgressCallback(total_files_num) - - def checkout(self, target=None, with_deps=False, force=False, recursive=False): from dvc.stage import StageFileDoesNotExistError, StageFileBadNameError @@ -44,7 +37,7 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False): with self.state: _cleanup_unused_links(self, all_stages) - progress_callback = get_progress_callback(stages) + pbar = Tqdm(total=get_all_files_numbers(stages)) for stage in stages: if stage.locked: @@ -53,6 +46,5 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False): " not going to be checked out.".format(path=stage.relpath) ) - stage.checkout(force=force, progress_callback=progress_callback) - if progress_callback: - progress_callback.finish("Checkout finished!") + stage.checkout(force=force, progress_callback=pbar.update_desc) + pbar.close()