From 1d617459e48270f3359ff7b151bd7a27855f903f Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <casper.dcl@physics.org>
Date: Mon, 12 Aug 2019 00:05:31 +0100
Subject: [PATCH] ensure nested bars are cleared away

---
 dvc/progress.py              | 19 +++++++++++++++++++
 dvc/remote/base.py           |  9 +++++----
 dvc/remote/local/__init__.py |  5 ++---
 dvc/remote/ssh/__init__.py   |  4 ++--
 4 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/dvc/progress.py b/dvc/progress.py
index 84ea9f4ed..b39eaf890 100644
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -2,6 +2,25 @@
 from __future__ import print_function
 import logging
 from tqdm import tqdm
+from concurrent.futures import ThreadPoolExecutor
+
+
+class TqdmThreadPoolExecutor(ThreadPoolExecutor):
+    """
+    Ensure worker progressbars are cleared away properly.
+    """
+
+    def __enter__(self):
+        """
+        Creates a blank initial dummy progress bar so that workers are forced to
+        create "nested" bars.
+        """
+        self.blank_bar = Tqdm(bar_format="Multi-Threaded:", leave=False)
+        super(TqdmThreadPoolExecutor, self).__enter__(self)
+
+    def __exit__(self, *a, **k):
+        super(TqdmThreadPoolExecutor, self).__exit__(*a, **k)
+        self.blank_bar.close()
 
 
 class Tqdm(tqdm):
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
index c2b6b0d50..9a2aeea35 100644
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -10,7 +10,6 @@ import tempfile
 import itertools
 from operator import itemgetter
 from multiprocessing import cpu_count
-from concurrent.futures import ThreadPoolExecutor
 import functools
 
 import dvc.prompt as prompt
@@ -20,7 +19,7 @@ from dvc.exceptions import (
     ConfirmRemoveError,
     DvcIgnoreInCollectedDirError,
 )
-from dvc.progress import Tqdm
+from dvc.progress import Tqdm, TqdmThreadPoolExecutor
 from dvc.utils import (
     LARGE_DIR_SIZE,
     tmp_fname,
@@ -151,7 +150,9 @@ class RemoteBASE(object):
 
     def _calculate_checksums(self, file_infos):
         file_infos = list(file_infos)
-        with ThreadPoolExecutor(max_workers=self.checksum_jobs) as executor:
+        with TqdmThreadPoolExecutor(
+            max_workers=self.checksum_jobs
+        ) as executor:
             tasks = executor.map(self.get_file_checksum, file_infos)
 
             if len(file_infos) > LARGE_DIR_SIZE:
@@ -635,7 +636,7 @@ class RemoteBASE(object):
                     self.batch_exists, callback=pbar.update_desc
                 )
 
-                with ThreadPoolExecutor(
+                with TqdmThreadPoolExecutor(
                     max_workers=jobs or self.JOBS
                 ) as executor:
                     path_infos = [
diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py
index 7c50a6435..0572e5945 100644
--- a/dvc/remote/local/__init__.py
+++ b/dvc/remote/local/__init__.py
@@ -35,8 +35,7 @@ from dvc.utils import (
 )
 from dvc.config import Config
 from dvc.exceptions import DvcException
-from dvc.progress import Tqdm
-from concurrent.futures import ThreadPoolExecutor
+from dvc.progress import Tqdm, TqdmThreadPoolExecutor
 
 from dvc.path_info import PathInfo
 
@@ -406,7 +405,7 @@ class RemoteLOCAL(RemoteBASE):
             return 0
 
         if jobs > 1:
-            with ThreadPoolExecutor(max_workers=jobs) as executor:
+            with TqdmThreadPoolExecutor(max_workers=jobs) as executor:
                 fails = sum(executor.map(func, *plans))
         else:
             fails = sum(map(func, *plans))
diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index 993ad563f..c817eabaa 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/dvc/remote/ssh/__init__.py
@@ -5,7 +5,6 @@ import getpass
 import logging
 import itertools
 import errno
-from concurrent.futures import ThreadPoolExecutor
 import threading
 
 try:
@@ -20,6 +19,7 @@ from dvc.utils.compat import urlparse, StringIO
 from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
 from dvc.remote.pool import get_connection
+from dvc.progress import TqdmThreadPoolExecutor
 
 from .connection import SSHConnection
 
@@ -167,7 +167,7 @@ class RemoteSSH(RemoteBASE):
             channels = ssh.open_max_sftp_channels()
             max_workers = len(channels)
 
-            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            with TqdmThreadPoolExecutor(max_workers=max_workers) as executor:
                 paths = [path_info.path for path_info in path_infos]
                 chunks = to_chunks(paths, num_chunks=max_workers)
                 chunks_and_channels = zip(chunks, channels)