Repository: iterative/dvc

Issue:
Use progress bar for dvc pull/checkout
As per discussion here:

https://discordapp.com/channels/485586884165107732/485586884165107734/547833091264086067

One of the suggestions is to use npm-style progress bar instead of printing the files one by one. Something similar to this: 

https://terminalizer.com/view/1f4a8cae675



Discussion:
Related https://github.com/iterative/dvc/issues/1430 
@efiop, @shcheklein Would we rather have our own progress bar, or can we consider 3rd party, like https://github.com/tqdm/tqdm ? It has MIT license, so I guess it would be ok to use it.
Don't see any reasons to do our own library. Looks like a great library, @pared !
@pared , I like `tqdm` too; the only reason I find is that every library that we add to `dvc` would need to be included in the packaged version.

By the way, `python-tqdm` is already on the `community` package repository of arch linux, it makes me think that is frequently used as a dependency.
We did try tqdm in the early stages of dvc, but couldn't make it work the way we wanted with multithreading and stuff. But maybe it will fit into the current dvc, there have been so many changes since then after all 🙂 
I forgot that we already have progress bar. I assume due to resons mentioned by @efiop. In that case I am not sure that introducing new dependency is good idea. We should probably stay with our current progress bar, or try to apply tqdm everywhere. In 2 case I think we should treat it as separate task. What do you guys think? 
Also there is one more thing to discuss before I do something:
I do believe that we could implement this feature by using `dvc.progress.progress` in `dvc.repo.checkout`. However there exists problem what to do with logged information, that is logged "below".  To mention those that I found so far:
-  `logger.info("Linking directory '{}'.".format(dir_relpath))` in `remote.local.do_checkout` 
- `"Stage '{}' didn't change.".format(self.relpath)` in `dvc.stage`

Those infos will break display of progress bar, and result partial progress bar updates mixed with logged information. How would we like to process from here? 
I see here two options:
1. Change info to debug, and don't care about ugly look in debug.
2. Make logger log to file somewhere under .dvc dir and keep minimalistic progress in stdout.

What do you guys think?
@pared I agree, migrating to tqdm is a separate task, unless our current progress bar is not suitable for the current task and tqdm is.

Have you tested that though? If recent progress bar changes didn't break anything, it should print info() cleanly if we are during progress bar print and start progress bar over. It is hard to explain in words, need to try and see :) Basically, it looks like progress bar is a last dynamic line in your terminal and all info()'s and debug()'s and etc are printed above it. 

Patch:
diff --git a/dvc/logger.py b/dvc/logger.py
--- a/dvc/logger.py
+++ b/dvc/logger.py
@@ -22,6 +22,10 @@ def info(message):
     logger.info(message)
 
 
+# Do not annotate debug with progress_aware. Why? Debug messages are called
+# in analytics.py. Trigger introduced by @progress_aware may run even
+# if log_level != DEBUG (look at Progress.print) Which might result in
+# messy terminal after any command (because analytics run in separate thread)
 def debug(message):
     """Prints a debug message."""
     prefix = colorize("Debug", color="blue")
diff --git a/dvc/output/base.py b/dvc/output/base.py
--- a/dvc/output/base.py
+++ b/dvc/output/base.py
@@ -197,12 +197,12 @@ def verify_metric(self):
     def download(self, to_info, resume=False):
         self.remote.download([self.path_info], [to_info], resume=resume)
 
-    def checkout(self, force=False):
+    def checkout(self, force=False, progress_callback=None):
         if not self.use_cache:
             return
 
         getattr(self.repo.cache, self.scheme).checkout(
-            self.path_info, self.info, force=force
+            output=self, force=force, progress_callback=progress_callback
         )
 
     def remove(self, ignore_remove=False):
@@ -225,3 +225,9 @@ def move(self, out):
 
         if self.scheme == "local" and self.use_cache and self.is_local:
             self.repo.scm.ignore(self.path)
+
+    def get_files_number(self):
+        if self.use_cache and self.checksum:
+            return 1
+
+        return 0
diff --git a/dvc/output/local.py b/dvc/output/local.py
--- a/dvc/output/local.py
+++ b/dvc/output/local.py
@@ -33,6 +33,8 @@ def __init__(
 
         self.path_info = {"scheme": "local", "path": p}
 
+        self._dir_cache = {}
+
     def __str__(self):
         return self.rel_path
 
@@ -69,6 +71,10 @@ def stage_path(self):
     def cache(self):
         return self.repo.cache.local.get(self.checksum)
 
+    @property
+    def is_dir_cache(self):
+        return self.repo.cache.local.is_dir_cache(self.checksum)
+
     def dumpd(self):
         ret = super(OutputLOCAL, self).dumpd()
         if self.is_local:
@@ -133,3 +139,20 @@ def save(self):
                 self.repo.scm.ignore(self.path)
 
         self.info = self.remote.save_info(self.path_info)
+
+    @property
+    def dir_cache(self):
+        if self.checksum not in self._dir_cache.keys():
+            self._dir_cache[
+                self.checksum
+            ] = self.repo.cache.local.load_dir_cache(self.checksum)
+
+        return self._dir_cache[self.checksum]
+
+    def get_files_number(self):
+        if self.cache is None:
+            return 0
+
+        if self.is_dir_cache:
+            return len(self.dir_cache)
+        return 1
diff --git a/dvc/progress.py b/dvc/progress.py
--- a/dvc/progress.py
+++ b/dvc/progress.py
@@ -8,6 +8,8 @@
 import sys
 import threading
 
+CLEARLINE_PATTERN = "\r\x1b[K"
+
 
 class Progress(object):
     """
@@ -30,12 +32,12 @@ def is_finished(self):
         """Returns if all targets have finished."""
         return self._n_total == self._n_finished
 
-    def _clearln(self):
-        self._print("\r\x1b[K", end="")
+    def clearln(self):
+        self.print(CLEARLINE_PATTERN, end="")
 
     def _writeln(self, line):
-        self._clearln()
-        self._print(line, end="")
+        self.clearln()
+        self.print(line, end="")
         sys.stdout.flush()
 
     def reset(self):
@@ -70,9 +72,9 @@ def finish_target(self, name):
             pbar = self._bar(name, 100, 100)
 
             if sys.stdout.isatty():
-                self._clearln()
+                self.clearln()
 
-            self._print(pbar)
+            self.print(pbar)
 
             self._n_finished += 1
             self._line = None
@@ -104,7 +106,7 @@ def _bar(self, target_name, current, total):
         return num + pbar + percent + target_name
 
     @staticmethod
-    def _print(*args, **kwargs):
+    def print(*args, **kwargs):
         import dvc.logger as logger
 
         if logger.is_quiet():
@@ -115,7 +117,7 @@ def _print(*args, **kwargs):
     def __enter__(self):
         self._lock.acquire(True)
         if self._line is not None:
-            self._clearln()
+            self.clearln()
 
     def __exit__(self, typ, value, tbck):
         if self._line is not None:
@@ -130,11 +132,25 @@ def progress_aware(f):
     @wraps(f)
     def wrapper(*args, **kwargs):
         if not progress.is_finished:
-            progress._print()
+            progress.print()
+        progress.clearln()
 
         return f(*args, **kwargs)
 
     return wrapper
 
 
+class ProgressCallback(object):
+    def __init__(self, total):
+        self.total = total
+        self.current = 0
+
+    def update(self, name, progress_to_add=1):
+        self.current += progress_to_add
+        progress.update_target(name, self.current, self.total)
+
+    def finish(self, name):
+        progress.finish_target(name)
+
+
 progress = Progress()  # pylint: disable=invalid-name
diff --git a/dvc/remote/base.py b/dvc/remote/base.py
--- a/dvc/remote/base.py
+++ b/dvc/remote/base.py
@@ -322,7 +322,10 @@ def safe_remove(self, path_info, force=False):
 
         self.remove(path_info)
 
-    def do_checkout(self, path_info, checksum, force=False):
+    def do_checkout(self, output, force=False, progress_callback=None):
+        path_info = output.path_info
+        checksum = output.info.get(self.PARAM_CHECKSUM)
+
         if self.exists(path_info):
             msg = "data '{}' exists. Removing before checkout."
             logger.warning(msg.format(str(path_info)))
@@ -331,29 +334,31 @@ def do_checkout(self, path_info, checksum, force=False):
         from_info = self.checksum_to_path_info(checksum)
         self.copy(from_info, path_info)
 
-    def checkout(self, path_info, checksum_info, force=False):
-        scheme = path_info["scheme"]
+    def checkout(self, output, force=False, progress_callback=None):
+        scheme = output.path_info["scheme"]
         if scheme not in ["", "local"] and scheme != self.scheme:
             raise NotImplementedError
 
-        checksum = checksum_info.get(self.PARAM_CHECKSUM)
+        checksum = output.info.get(self.PARAM_CHECKSUM)
         if not checksum:
             msg = "No checksum info for '{}'."
-            logger.info(msg.format(str(path_info)))
+            logger.debug(msg.format(str(output.path_info)))
             return
 
-        if not self.changed(path_info, checksum_info):
+        if not self.changed(output.path_info, output.info):
             msg = "Data '{}' didn't change."
-            logger.info(msg.format(str(path_info)))
+            logger.debug(msg.format(str(output.path_info)))
             return
 
         if self.changed_cache(checksum):
             msg = "Cache '{}' not found. File '{}' won't be created."
-            logger.warning(msg.format(checksum, str(path_info)))
-            self.safe_remove(path_info, force=force)
+            logger.warning(msg.format(checksum, str(output.path_info)))
+            self.safe_remove(output.path_info, force=force)
             return
 
         msg = "Checking out '{}' with cache '{}'."
-        logger.info(msg.format(str(path_info), checksum))
+        logger.debug(msg.format(str(output.path_info), checksum))
 
-        self.do_checkout(path_info, checksum, force=force)
+        self.do_checkout(
+            output, force=force, progress_callback=progress_callback
+        )
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -291,17 +291,22 @@ def dump_dir_cache(self, md5, dir_info):
     def is_dir_cache(cls, cache):
         return cache.endswith(cls.MD5_DIR_SUFFIX)
 
-    def do_checkout(self, path_info, checksum, force=False):
+    def do_checkout(self, output, force=False, progress_callback=None):
+        path_info = output.path_info
+        checksum = output.info.get(self.PARAM_CHECKSUM)
+
         path = path_info["path"]
         md5 = checksum
         cache = self.get(md5)
 
-        if not self.is_dir_cache(cache):
+        if not output.is_dir_cache:
             if os.path.exists(path):
                 self.safe_remove(path_info, force=force)
 
             self.link(cache, path)
             self.state.update_link(path)
+            if progress_callback:
+                progress_callback.update(os.path.relpath(path))
             return
 
         # Create dir separately so that dir is created
@@ -309,12 +314,10 @@ def do_checkout(self, path_info, checksum, force=False):
         if not os.path.exists(path):
             os.makedirs(path)
 
-        dir_info = self.load_dir_cache(md5)
+        dir_info = output.dir_cache
         dir_relpath = os.path.relpath(path)
-        dir_size = len(dir_info)
-        bar = dir_size > LARGE_DIR_SIZE
 
-        logger.info("Linking directory '{}'.".format(dir_relpath))
+        logger.debug("Linking directory '{}'.".format(dir_relpath))
 
         for processed, entry in enumerate(dir_info):
             relpath = entry[self.PARAM_RELPATH]
@@ -332,16 +335,13 @@ def do_checkout(self, path_info, checksum, force=False):
 
                 self.link(c, p)
 
-            if bar:
-                progress.update_target(dir_relpath, processed, dir_size)
+            if progress_callback:
+                progress_callback.update(os.path.relpath(p))
 
         self._discard_working_directory_changes(path, dir_info, force=force)
 
         self.state.update_link(path)
 
-        if bar:
-            progress.finish_target(dir_relpath)
-
     def already_cached(self, path_info):
         assert path_info["scheme"] in ["", "local"]
 
diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py
--- a/dvc/repo/checkout.py
+++ b/dvc/repo/checkout.py
@@ -2,6 +2,7 @@
 
 import dvc.logger as logger
 from dvc.exceptions import DvcException
+from dvc.progress import ProgressCallback
 
 
 def _cleanup_unused_links(self, all_stages):
@@ -12,6 +13,17 @@ def _cleanup_unused_links(self, all_stages):
     self.state.remove_unused_links(used)
 
 
+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
 
@@ -29,6 +41,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)
 
         for stage in stages:
             if stage.locked:
@@ -37,4 +50,6 @@ 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)
+            stage.checkout(force=force, progress_callback=progress_callback)
+        if progress_callback:
+            progress_callback.finish("Checkout finished!")
diff --git a/dvc/stage.py b/dvc/stage.py
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -739,9 +739,9 @@ def check_missing_outputs(self):
         if paths:
             raise MissingDataSource(paths)
 
-    def checkout(self, force=False):
+    def checkout(self, force=False, progress_callback=None):
         for out in self.outs:
-            out.checkout(force=force)
+            out.checkout(force=force, progress_callback=progress_callback)
 
     @staticmethod
     def _status(entries):
@@ -784,3 +784,6 @@ def _already_cached(self):
                 for out in self.outs
             )
         )
+
+    def get_all_files_number(self):
+        return sum(out.get_files_number() for out in self.outs)


Test patch:
diff --git a/tests/test_checkout.py b/tests/test_checkout.py
--- a/tests/test_checkout.py
+++ b/tests/test_checkout.py
@@ -1,11 +1,15 @@
 import os
-import mock
+import sys
+import re
+
 import yaml
 import shutil
 import filecmp
 import collections
 
+from dvc.logger import logger
 from dvc.main import main
+from dvc import progress
 from dvc.repo import Repo as DvcRepo
 from dvc.system import System
 from tests.basic_env import TestDvc
@@ -15,6 +19,7 @@
 from dvc.exceptions import DvcException, ConfirmRemoveError
 
 from mock import patch
+from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 
 class TestCheckout(TestRepro):
@@ -355,8 +360,8 @@ def test(self):
 
 
 class TestCheckoutHook(TestDvc):
-    @mock.patch("sys.stdout.isatty", return_value=True)
-    @mock.patch("dvc.prompt.input", side_effect=EOFError)
+    @patch("sys.stdout.isatty", return_value=True)
+    @patch("dvc.prompt.input", side_effect=EOFError)
     def test(self, mock_input, mock_isatty):
         """ Test that dvc checkout handles EOFError gracefully, which is what
         it will experience when running in a git hook.
@@ -402,3 +407,108 @@ def test(self):
                     )
                 ),
             )
+
+
+class TestCheckoutShouldHaveSelfClearingProgressBar(TestDvc):
+    def setUp(self):
+        super(TestCheckoutShouldHaveSelfClearingProgressBar, self).setUp()
+        self._prepare_repo()
+
+    def test(self):
+        self._checkout_and_intercept_std_output()
+
+        stdout_calls = self.stdout_mock.method_calls
+        write_calls = self.filter_out_non_write_calls(stdout_calls)
+        write_calls = self.filter_out_empty_write_calls(write_calls)
+        self.write_args = [w_c[1][0] for w_c in write_calls]
+
+        progress_bars = self.get_progress_bars()
+        update_bars = progress_bars[:-1]
+        finish_bar = progress_bars[-1]
+
+        self.assertEqual(4, len(update_bars))
+        self.assertRegexpMatches(progress_bars[0], ".*\\[#{7} {23}\\] 25%.*")
+        self.assertRegexpMatches(progress_bars[1], ".*\\[#{15} {15}\\] 50%.*")
+        self.assertRegexpMatches(progress_bars[2], ".*\\[#{22} {8}\\] 75%.*")
+        self.assertRegexpMatches(progress_bars[3], ".*\\[#{30}\\] 100%.*")
+
+        self.assertCaretReturnFollowsEach(update_bars)
+        self.assertNewLineFollows(finish_bar)
+
+        self.assertAnyEndsWith(update_bars, self.FOO)
+        self.assertAnyEndsWith(update_bars, self.BAR)
+        self.assertAnyEndsWith(update_bars, self.DATA)
+        self.assertAnyEndsWith(update_bars, self.DATA_SUB)
+
+        self.assertTrue(finish_bar.endswith("Checkout finished!"))
+
+    def filter_out_empty_write_calls(self, calls):
+        def is_not_empty_write(call):
+            assert call[0] == "write"
+            return call[1][0] != ""
+
+        return list(filter(is_not_empty_write, calls))
+
+    def filter_out_non_write_calls(self, calls):
+        def is_write_call(call):
+            return call[0] == "write"
+
+        return list(filter(is_write_call, calls))
+
+    def _prepare_repo(self):
+        storage = self.mkdtemp()
+
+        ret = main(["remote", "add", "-d", "myremote", storage])
+        self.assertEqual(0, ret)
+
+        ret = main(["add", self.DATA_DIR])
+        self.assertEqual(0, ret)
+
+        ret = main(["add", self.FOO])
+        self.assertEqual(0, ret)
+
+        ret = main(["add", self.BAR])
+        self.assertEqual(0, ret)
+
+        ret = main(["push"])
+        self.assertEqual(0, ret)
+
+        shutil.rmtree(self.DATA_DIR)
+        os.unlink(self.FOO)
+        os.unlink(self.BAR)
+
+    def _checkout_and_intercept_std_output(self):
+        with MockLoggerHandlers(
+            logger
+        ), ConsoleFontColorsRemover(), patch.object(
+            sys, "stdout"
+        ) as stdout_mock:
+            self.stdout_mock = logger.handlers[0].stream = stdout_mock
+
+            ret = main(["checkout"])
+            self.assertEqual(0, ret)
+
+    def get_progress_bars(self):
+        pattern = re.compile(".*\\[.{30}\\].*%.*")
+        return [arg for arg in self.write_args if pattern.match(arg)]
+
+    def assertCaretReturnFollowsEach(self, update_bars):
+        for update_bar in update_bars:
+
+            self.assertIn(update_bar, self.write_args)
+
+            for index, arg in enumerate(self.write_args):
+                if arg == update_bar:
+                    self.assertEqual(
+                        progress.CLEARLINE_PATTERN, self.write_args[index + 1]
+                    )
+
+    def assertNewLineFollows(self, finish_bar):
+        self.assertIn(finish_bar, self.write_args)
+
+        for index, arg in enumerate(self.write_args):
+            if arg == finish_bar:
+                self.assertEqual("\n", self.write_args[index + 1])
+
+    def assertAnyEndsWith(self, update_bars, name):
+        self.assertTrue(any(ub for ub in update_bars if ub.endswith(name)))
diff --git a/tests/unit/output/test_local.py b/tests/unit/output/test_local.py
--- a/tests/unit/output/test_local.py
+++ b/tests/unit/output/test_local.py
@@ -1,4 +1,4 @@
-import mock
+from mock import patch
 
 from dvc.stage import Stage
 from dvc.output import OutputLOCAL
@@ -16,6 +16,33 @@ def _get_output(self):
 
     def test_save_missing(self):
         o = self._get_output()
-        with mock.patch.object(o.remote, "exists", return_value=False):
+        with patch.object(o.remote, "exists", return_value=False):
             with self.assertRaises(o.DoesNotExistError):
                 o.save()
+
+
+class TestGetFilesNumber(TestDvc):
+    def _get_output(self):
+        stage = Stage(self.dvc)
+        return OutputLOCAL(stage, "path")
+
+    @patch.object(OutputLOCAL, "cache", None)
+    def test_return_0_on_no_cache(self):
+        o = self._get_output()
+
+        self.assertEqual(0, o.get_files_number())
+
+    @patch.object(OutputLOCAL, "cache", "cache")
+    @patch.object(OutputLOCAL, "is_dir_cache", True)
+    @patch.object(OutputLOCAL, "dir_cache", [{"md5": "asdf"}, {"md5": "qwe"}])
+    def test_return_mutiple_for_dir(self):
+        o = self._get_output()
+
+        self.assertEqual(2, o.get_files_number())
+
+    @patch.object(OutputLOCAL, "cache", "cache")
+    @patch.object(OutputLOCAL, "is_dir_cache", False)
+    def test_return_1_on_single_file_cache(self):
+        o = self._get_output()
+
+        self.assertEqual(1, o.get_files_number())
diff --git a/tests/unit/test_progress.py b/tests/unit/test_progress.py
--- a/tests/unit/test_progress.py
+++ b/tests/unit/test_progress.py
@@ -6,7 +6,7 @@
 
 class TestProgressAware(TestCase):
     @mock.patch("sys.stdout.isatty", return_value=True)
-    @mock.patch("dvc.progress.Progress._print")
+    @mock.patch("dvc.progress.Progress.print")
     def test(self, mock_print, _):
         # progress is a global object, can be shared between tests when
         # run in multi-threading environment with pytest
@@ -14,10 +14,14 @@ def test(self, mock_print, _):
         function = progress_aware(lambda: None)
 
         function()
-        mock_print.assert_not_called()
+        # first - called once for progress.clearln()
+        mock_print.assert_called_once()
 
         progress.update_target("testing", 0, 100)
         function()
-        mock_print.assert_called()
+        # second - progress.clearln() on refresh on update_target
+        # third - progress.print on refresh on update_target
+        # fourth - progress.clearln()
+        self.assertEqual(4, mock_print.call_count)
 
         progress.finish_target("testing")