From e1eab9405cb10c992d4bde5a25dd469f02488d6f Mon Sep 17 00:00:00 2001
From: pawel <pawelredzynski@gmail.com>
Date: Fri, 27 Dec 2019 21:54:47 +0100
Subject: [PATCH] tree: check that repo.tree is CleanTree when using walk_files

---
 dvc/remote/local.py       | 5 +++++
 dvc/utils/fs.py           | 4 ++++
 tests/func/test_ignore.py | 7 ++++---
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/dvc/remote/local.py b/dvc/remote/local.py
index 81b63d061..7a68ef824 100644
--- a/dvc/remote/local.py
+++ b/dvc/remote/local.py
@@ -13,6 +13,7 @@ from dvc.config import Config
 from dvc.exceptions import DownloadError
 from dvc.exceptions import DvcException
 from dvc.exceptions import UploadError
+from dvc.ignore import CleanTree
 from dvc.path_info import PathInfo
 from dvc.progress import Tqdm
 from dvc.remote.base import RemoteBASE
@@ -138,6 +139,8 @@ class RemoteLOCAL(RemoteBASE):
         return os.path.getsize(fspath_py35(path_info))
 
     def walk_files(self, path_info):
+        assert isinstance(self.repo.tree, CleanTree)
+
         for fname in self.repo.tree.walk_files(path_info):
             yield PathInfo(fname)
 
@@ -427,6 +430,8 @@ class RemoteLOCAL(RemoteBASE):
         os.chmod(path, os.stat(path).st_mode | stat.S_IWRITE)
 
     def _unprotect_dir(self, path):
+        assert isinstance(self.repo.tree, CleanTree)
+
         for fname in self.repo.tree.walk_files(path):
             RemoteLOCAL._unprotect_file(fname)
 
diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py
index ebb4b6cbc..4f1a8013a 100644
--- a/dvc/utils/fs.py
+++ b/dvc/utils/fs.py
@@ -35,6 +35,10 @@ def get_inode(path):
 
 
 def get_mtime_and_size(path, tree):
+    from dvc.ignore import CleanTree
+
+    assert isinstance(tree, CleanTree)
+
     if os.path.isdir(fspath_py35(path)):
         size = 0
         files_mtimes = {}
diff --git a/tests/func/test_ignore.py b/tests/func/test_ignore.py
index 22fef716e..68137300a 100644
--- a/tests/func/test_ignore.py
+++ b/tests/func/test_ignore.py
@@ -116,10 +116,11 @@ def test_ignore_on_branch(tmp_dir, scm, dvc):
     assert _files_set(".", dvc.tree) == {"./foo", "./bar"}
 
     tree = CleanTree(scm.get_tree("branch"))
-    files_on_branch = {
-        f for r, ds, fs in tree.walk(tree.tree_root) for f in fs
+
+    assert _files_set(".", tree) == {
+        os.path.join(tree.tree_root, DvcIgnore.DVCIGNORE_FILE),
+        os.path.join(tree.tree_root, "bar"),
     }
-    assert "foo" not in files_on_branch
 
 
 def _files_set(root, tree):