From 5e92a079b87af3efffde6697fcacb9163eff1fdd Mon Sep 17 00:00:00 2001
From: Alexander Schepanovski <suor.web@gmail.com>
Date: Fri, 3 Jan 2020 21:46:05 +0700
Subject: [PATCH] dvc: autowrap repo tree into CleanTree

---
 dvc/repo/__init__.py      |  4 ++--
 dvc/repo/brancher.py      |  5 ++---
 dvc/repo/diff.py          | 14 +++++---------
 tests/func/test_ignore.py | 11 +++++------
 4 files changed, 14 insertions(+), 20 deletions(-)

diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
index 3c35683be..2376acb6a 100644
--- a/dvc/repo/__init__.py
+++ b/dvc/repo/__init__.py
@@ -83,7 +83,7 @@ class Repo(object):
 
         self.scm = SCM(self.root_dir)
 
-        self.tree = CleanTree(WorkingTree(self.root_dir))
+        self.tree = WorkingTree(self.root_dir)
 
         self.tmp_dir = os.path.join(self.dvc_dir, "tmp")
         makedirs(self.tmp_dir, exist_ok=True)
@@ -120,7 +120,7 @@ class Repo(object):
 
     @tree.setter
     def tree(self, tree):
-        self._tree = tree
+        self._tree = tree if isinstance(tree, CleanTree) else CleanTree(tree)
         # Our graph cache is no longer valid, as it was based on the previous
         # tree.
         self._reset()
diff --git a/dvc/repo/brancher.py b/dvc/repo/brancher.py
index ca6dfe84f..8b096726c 100644
--- a/dvc/repo/brancher.py
+++ b/dvc/repo/brancher.py
@@ -1,6 +1,5 @@
 from funcy import group_by
 
-from dvc.ignore import CleanTree
 from dvc.scm.tree import WorkingTree
 
 
@@ -30,7 +29,7 @@ def brancher(  # noqa: E302
 
     scm = self.scm
 
-    self.tree = CleanTree(WorkingTree(self.root_dir))
+    self.tree = WorkingTree(self.root_dir)
     yield "working tree"
 
     if all_commits:
@@ -44,7 +43,7 @@ def brancher(  # noqa: E302
 
     try:
         for sha, names in group_by(scm.resolve_rev, revs).items():
-            self.tree = CleanTree(scm.get_tree(sha))
+            self.tree = scm.get_tree(sha)
             yield ", ".join(names)
     finally:
         self.tree = saved_tree
diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py
index f4737da72..165b286b9 100644
--- a/dvc/repo/diff.py
+++ b/dvc/repo/diff.py
@@ -1,15 +1,11 @@
 import os
 from errno import ENOENT
 
-import dvc.logger as logger
-from dvc.ignore import CleanTree
-from . import locked
+from dvc import logger
 from dvc.scm.base import FileNotInCommitError
-from dvc.scm.git import DIFF_A_REF
-from dvc.scm.git import DIFF_A_TREE
-from dvc.scm.git import DIFF_B_REF
-from dvc.scm.git import DIFF_B_TREE
+from dvc.scm.git import DIFF_A_REF, DIFF_B_REF, DIFF_A_TREE, DIFF_B_TREE
 from dvc.scm.git import DIFF_EQUAL
+from . import locked
 
 
 DIFF_TARGET = "target"
@@ -136,9 +132,9 @@ def _is_dir(path, a_outs, b_outs):
 
 
 def _get_diff_outs(self, diff_dct):
-    self.tree = CleanTree(diff_dct[DIFF_A_TREE])
+    self.tree = diff_dct[DIFF_A_TREE]
     a_outs = {str(out): out for st in self.stages for out in st.outs}
-    self.tree = CleanTree(diff_dct[DIFF_B_TREE])
+    self.tree = diff_dct[DIFF_B_TREE]
     b_outs = {str(out): out for st in self.stages for out in st.outs}
     outs_paths = set(a_outs.keys())
     outs_paths.update(b_outs.keys())
diff --git a/tests/func/test_ignore.py b/tests/func/test_ignore.py
index ff722d5db..2ba0d485a 100644
--- a/tests/func/test_ignore.py
+++ b/tests/func/test_ignore.py
@@ -4,7 +4,7 @@ import shutil
 import pytest
 
 from dvc.exceptions import DvcIgnoreInCollectedDirError
-from dvc.ignore import DvcIgnore, DvcIgnoreDirs, DvcIgnorePatterns, CleanTree
+from dvc.ignore import DvcIgnore, DvcIgnoreDirs, DvcIgnorePatterns
 from dvc.scm.tree import WorkingTree
 from dvc.utils import relpath
 from dvc.compat import fspath_py35, fspath
@@ -113,11 +113,10 @@ def test_ignore_on_branch(tmp_dir, scm, dvc):
     scm.checkout("master")
     assert _files_set(".", dvc.tree) == {"./foo", "./bar"}
 
-    tree = CleanTree(scm.get_tree("branch"))
-
-    assert _files_set(".", tree) == {
-        to_posixpath(os.path.join(tree.tree_root, DvcIgnore.DVCIGNORE_FILE)),
-        to_posixpath(os.path.join(tree.tree_root, "bar")),
+    dvc.tree = scm.get_tree("branch")
+    assert _files_set(".", dvc.tree) == {
+        to_posixpath(os.path.join(dvc.root_dir, DvcIgnore.DVCIGNORE_FILE)),
+        to_posixpath(os.path.join(dvc.root_dir, "bar")),
     }