diff --git a/dvc/progress.py b/dvc/progress.py index ea0faa2ca..b25f07d43 100644 --- a/dvc/progress.py +++ b/dvc/progress.py @@ -126,6 +126,16 @@ class Progress(object): self.refresh() self._lock.release() + def __call__(self, seq, name="", total=None): + if total is None: + total = len(seq) + + self.update_target(name, 0, total) + for done, item in enumerate(seq, start=1): + yield item + self.update_target(name, done, total) + self.finish_target(name) + class ProgressCallback(object): def __init__(self, total): diff --git a/dvc/remote/base.py b/dvc/remote/base.py index fa85fa170..870a51a55 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -15,7 +15,8 @@ from multiprocessing import cpu_count import dvc.prompt as prompt from dvc.config import Config from dvc.exceptions import DvcException, ConfirmRemoveError -from dvc.utils import tmp_fname +from dvc.progress import progress +from dvc.utils import LARGE_DIR_SIZE, tmp_fname from dvc.state import StateBase @@ -161,6 +162,15 @@ class RemoteBase(object): p_info = path_info.copy() dpath = p_info["path"] for root, dirs, files in self.walk(path_info): + if len(files) > LARGE_DIR_SIZE: + msg = ( + "Computing md5 for a large directory {}. " + "This is only done once." + ) + relpath = self.ospath.relpath(root) + logger.info(msg.format(relpath)) + files = progress(files, name=relpath) + for fname in files: path = self.ospath.join(root, fname) p_info["path"] = path diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index 528a8ff06..71f5ca0ad 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -1,5 +1,6 @@ import os import configobj +from mock import patch from dvc.main import main from dvc.command.remote import CmdRemoteAdd @@ -169,3 +170,16 @@ class TestRemoteShouldHandleUppercaseRemoteName(TestDvc): ret = main(["push", "-r", self.upper_case_remote_name]) self.assertEqual(ret, 0) + + +def test_large_dir_progress(repo_dir, dvc): + from dvc.utils import LARGE_DIR_SIZE + from dvc.progress import progress + + # Create a "large dir" + for i in range(LARGE_DIR_SIZE + 1): + repo_dir.create(os.path.join("gen", "{}.txt".format(i)), str(i)) + + with patch.object(progress, "update_target") as update_target: + dvc.add("gen") + assert update_target.called