diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 224744d52..5b527806e 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -11,6 +11,7 @@ import itertools from operator import itemgetter from multiprocessing import cpu_count from concurrent.futures import ThreadPoolExecutor +import functools import dvc.prompt as prompt from dvc.config import Config @@ -628,20 +629,23 @@ class RemoteBASE(object): Returns: A list with checksums that were found in the remote """ - pbar = Tqdm(total=len(checksums)) - - def exists_with_progress(chunks): - return self.batch_exists(chunks, callback=pbar.update_desc) - if self.no_traverse and hasattr(self, "batch_exists"): - with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor: - path_infos = [self.checksum_to_path_info(x) for x in checksums] - chunks = to_chunks(path_infos, num_chunks=self.JOBS) - results = executor.map(exists_with_progress, chunks) - in_remote = itertools.chain.from_iterable(results) - ret = list(itertools.compress(checksums, in_remote)) - pbar.close() - return ret + with Tqdm(total=len(checksums)) as pbar: + exists_with_progress = functools.partial( + self.batch_exists, callback=pbar.update_desc + ) + + with ThreadPoolExecutor( + max_workers=jobs or self.JOBS + ) as executor: + path_infos = [ + self.checksum_to_path_info(x) for x in checksums + ] + chunks = to_chunks(path_infos, num_chunks=self.JOBS) + results = executor.map(exists_with_progress, chunks) + in_remote = itertools.chain.from_iterable(results) + ret = list(itertools.compress(checksums, in_remote)) + return ret return list(set(checksums) & set(self.all())) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 72eaf41d1..40f47c104 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -229,9 +229,12 @@ class RemoteS3(RemoteBASE): ) def _download(self, from_info, to_file, name=None, no_progress_bar=False): - total = self.s3.head_object( - Bucket=from_info.bucket, Key=from_info.path - )["ContentLength"] + if no_progress_bar: + total = None + else: + total = self.s3.head_object( + Bucket=from_info.bucket, Key=from_info.path + )["ContentLength"] with Tqdm( disable=no_progress_bar, total=total, diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py index 826a24a81..0ddbfc4d7 100644 --- a/dvc/repo/checkout.py +++ b/dvc/repo/checkout.py @@ -37,14 +37,14 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False): with self.state: _cleanup_unused_links(self, all_stages) - pbar = Tqdm(total=get_all_files_numbers(stages)) + with Tqdm(total=get_all_files_numbers(stages)) as pbar: + for stage in stages: + if stage.locked: + logger.warning( + "DVC-file '{path}' is locked. Its dependencies are" + " not going to be checked out.".format( + path=stage.relpath + ) + ) - for stage in stages: - if stage.locked: - logger.warning( - "DVC-file '{path}' is locked. Its dependencies are" - " not going to be checked out.".format(path=stage.relpath) - ) - - stage.checkout(force=force, progress_callback=pbar.update_desc) - pbar.close() + stage.checkout(force=force, progress_callback=pbar.update_desc)