Merge pull request #3018 from pared/2914_3

local-related methods should verify that WorkingTree is in use

Author: efiop

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,7 +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)
+        assert is_working_tree(self.repo.tree)
 
         for fname in self.repo.tree.walk_files(path):
             RemoteLOCAL._unprotect_file(fname)

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
+    )

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,11 +33,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)):
+        assert is_working_tree(tree)
+
         size = 0
         files_mtimes = {}
         for file_path in tree.walk_files(path):