misc minor restructuring

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()))
 

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,

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)