tree: introduce is_working_tree

iterative · Jan 3, 2020 · 1e1a9cf · 1e1a9cf
1 parent eb53c03
commit 1e1a9cf
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 12 deletions.
diff --git a/dvc/remote/local.py b/dvc/remote/local.py
@@ -11,7 +11,6 @@
 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
@@ -20,7 +19,7 @@
 from dvc.remote.base import STATUS_MISSING
 from dvc.remote.base import STATUS_NEW
 from dvc.scheme import Schemes
-from dvc.scm.tree import WorkingTree
+from dvc.scm.tree import is_working_tree
 from dvc.system import System
 from dvc.utils import copyfile
 from dvc.utils import file_md5
@@ -136,9 +135,7 @@ def getsize(path_info):
         return os.path.getsize(fspath_py35(path_info))
 
     def walk_files(self, path_info):
-        assert isinstance(self.repo.tree, CleanTree) and isinstance(
-            self.repo.tree.tree, WorkingTree
-        )
+        assert is_working_tree(self.repo.tree)
 
         for fname in self.repo.tree.walk_files(path_info):
             yield PathInfo(fname)
@@ -429,9 +426,7 @@ def _unprotect_file(path):
         os.chmod(path, os.stat(path).st_mode | stat.S_IWRITE)
 
     def _unprotect_dir(self, path):
-        assert isinstance(self.repo.tree, CleanTree) and isinstance(
-            self.repo.tree.tree, WorkingTree
-        )
+        assert is_working_tree(self.repo.tree)
 
         for fname in self.repo.tree.walk_files(path):
             RemoteLOCAL._unprotect_file(fname)

diff --git a/dvc/scm/tree.py b/dvc/scm/tree.py
@@ -74,3 +74,9 @@ def onerror(e):
             top, topdown=topdown, onerror=onerror
         ):
             yield os.path.normpath(root), dirs, files
+
+
+def is_working_tree(tree):
+    return isinstance(tree, WorkingTree) or isinstance(
+        getattr(tree, "tree", None), WorkingTree
+    )
diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py
@@ -8,6 +8,7 @@
 from shortuuid import uuid
 
 from dvc.exceptions import DvcException
+from dvc.scm.tree import is_working_tree
 from dvc.system import System
 from dvc.utils import dict_md5
 from dvc.utils import fspath
@@ -32,10 +33,7 @@ def get_inode(path):
 
 
 def get_mtime_and_size(path, tree):
-    from dvc.ignore import CleanTree
-    from dvc.scm.tree import WorkingTree
-
-    assert isinstance(tree, CleanTree) and isinstance(tree.tree, WorkingTree)
+    assert is_working_tree(tree)
 
     if os.path.isdir(fspath_py35(path)):
         size = 0