checkout: progress bar (#1713)

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

dvc/output/base.py

@@ -197,12 +197,12 @@ class OutputBase(object):
     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 @@ class OutputBase(object):
 
         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

dvc/output/local.py

@@ -33,6 +33,8 @@ class OutputLOCAL(OutputBase):
 
         self.path_info = {"scheme": "local", "path": p}
 
+        self._dir_cache = {}
+
     def __str__(self):
         return self.rel_path
 
@@ -69,6 +71,10 @@ class OutputLOCAL(OutputBase):
     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 @@ class OutputLOCAL(OutputBase):
                 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

dvc/progress.py

@@ -8,6 +8,8 @@ from dvc.utils.compat import str
 import sys
 import threading
 
+CLEARLINE_PATTERN = "\r\x1b[K"
+
 
 class Progress(object):
     """
@@ -30,12 +32,12 @@ class Progress(object):
         """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 @@ class Progress(object):
             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 @@ class Progress(object):
         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 @@ class Progress(object):
     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

dvc/remote/base.py

@@ -322,7 +322,10 @@ class RemoteBase(object):
 
         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 @@ class RemoteBase(object):
         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
+        )

dvc/remote/local.py

@@ -291,17 +291,22 @@ class RemoteLOCAL(RemoteBase):
     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 @@ class RemoteLOCAL(RemoteBase):
         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 @@ class RemoteLOCAL(RemoteBase):
 
                 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"]
 

dvc/repo/checkout.py

@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 
 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!")

dvc/stage.py

@@ -739,9 +739,9 @@ class Stage(object):
         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 @@ class Stage(object):
                 for out in self.outs
             )
         )
+
+    def get_all_files_number(self):
+        return sum(out.get_files_number() for out in self.outs)

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.remote.local import RemoteLOCAL
 from dvc.exceptions import DvcException, ConfirmRemoveError
 
 from mock import patch
+from tests.utils.logger import MockLoggerHandlers, ConsoleFontColorsRemover
 
 
 class TestCheckout(TestRepro):
@@ -355,8 +360,8 @@ class TestCheckoutDirectory(TestRepro):
 
 
 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 @@ class TestCheckoutSuggestGit(TestRepro):
                     )
                 ),
             )
+
+
+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)))

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 @@ class TestOutputLOCAL(TestDvc):
 
     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())

tests/unit/test_progress.py

@@ -6,7 +6,7 @@ from dvc.progress import progress, progress_aware
 
 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 @@ class TestProgressAware(TestCase):
         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")