diff --git a/dvc/logger.py b/dvc/logger.py index 1352e3d0a..d756c0d6f 100644 --- 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 index 67d0917a7..53dfe6f7e 100644 --- a/dvc/output/base.py +++ b/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 diff --git a/dvc/output/local.py b/dvc/output/local.py index de320554d..0bfe28d8e 100644 --- a/dvc/output/local.py +++ b/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 diff --git a/dvc/progress.py b/dvc/progress.py index 6245c14de..d101b7bf9 100644 --- a/dvc/progress.py +++ b/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 diff --git a/dvc/remote/base.py b/dvc/remote/base.py index a4965cba6..42f069e49 100644 --- a/dvc/remote/base.py +++ b/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 + ) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 44cdb06f0..9c4f35898 100644 --- a/dvc/remote/local.py +++ b/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"] diff --git a/dvc/repo/checkout.py b/dvc/repo/checkout.py index 2f5a74a07..21a350634 100644 --- a/dvc/repo/checkout.py +++ b/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!") diff --git a/dvc/stage.py b/dvc/stage.py index 3d4760058..c8b7bc491 100644 --- a/dvc/stage.py +++ b/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) diff --git a/tests/test_checkout.py b/tests/test_checkout.py index d1a2b5dc7..aa41b0430 100644 --- 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.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))) diff --git a/tests/unit/output/test_local.py b/tests/unit/output/test_local.py index 2a5895842..a9b74dc24 100644 --- 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 @@ 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()) diff --git a/tests/unit/test_progress.py b/tests/unit/test_progress.py index 94e40d79b..76e3ce9f8 100644 --- a/tests/unit/test_progress.py +++ b/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")