trees: handle tree.walk errors the same way as os.walk (#4059)

* trees: BaseTree.walk should handle errors like os.walk

* use onerror callback if exceptions should be logged or raised

* repo.ls: use onerror for tree.walk

dvc/ignore.py

@@ -229,8 +229,10 @@ class CleanTree(BaseTree):
                 return False
         return True
 
-    def walk(self, top, topdown=True):
-        for root, dirs, files in self.tree.walk(top, topdown):
+    def walk(self, top, topdown=True, onerror=None):
+        for root, dirs, files in self.tree.walk(
+            top, topdown=topdown, onerror=onerror
+        ):
             dirs[:], files[:] = self.dvcignore(
                 os.path.abspath(root), dirs, files
             )

dvc/repo/ls.py

@@ -50,13 +50,18 @@ def ls(
 def _ls(repo, path_info, recursive=None, dvc_only=False):
     from dvc.repo.tree import RepoTree
 
+    def onerror(exc):
+        raise exc
+
     # use our own RepoTree instance instead of repo.repo_tree since we do not
     # want fetch/stream enabled for ls
     tree = RepoTree(repo)
 
     ret = {}
     try:
-        for root, dirs, files in tree.walk(path_info.fspath, dvcfiles=True):
+        for root, dirs, files in tree.walk(
+            path_info.fspath, onerror=onerror, dvcfiles=True
+        ):
             for fname in files:
                 info = PathInfo(root) / fname
                 dvc = tree.isdvc(info)

dvc/repo/tree.py

@@ -162,16 +162,22 @@ class DvcTree(BaseTree):
         else:
             assert False
 
-    def walk(self, top, topdown=True, download_callback=None, **kwargs):
+    def walk(
+        self, top, topdown=True, onerror=None, download_callback=None, **kwargs
+    ):
         from pygtrie import Trie
 
         assert topdown
 
         if not self.exists(top):
-            raise FileNotFoundError
+            if onerror is not None:
+                onerror(FileNotFoundError(top))
+            return
 
         if not self.isdir(top):
-            raise NotADirectoryError
+            if onerror is not None:
+                onerror(NotADirectoryError(top))
+            return
 
         root = PathInfo(os.path.abspath(top))
         outs = self._find_outs(top, recursive=True, strict=False)
@@ -333,12 +339,14 @@ class RepoTree(BaseTree):
             elif dirname in repo_set:
                 yield from self._walk_one(repo_walk)
 
-    def walk(self, top, topdown=True, dvcfiles=False, **kwargs):
+    def walk(self, top, topdown=True, onerror=None, dvcfiles=False, **kwargs):
         """Walk and merge both DVC and repo trees.
 
         Args:
             top: path to walk from
             topdown: if True, tree will be walked from top down.
+            onerror: if set, onerror function will be called if an error
+                occurs (by default errors are ignored).
             dvcfiles: if True, dvcfiles will be included in the files list
                 for walked directories.
 
@@ -348,21 +356,27 @@ class RepoTree(BaseTree):
         assert topdown
 
         if not self.exists(top):
-            raise FileNotFoundError
+            if onerror is not None:
+                onerror(FileNotFoundError(top))
+            return
 
         if not self.isdir(top):
-            raise NotADirectoryError
+            if onerror is not None:
+                onerror(NotADirectoryError(top))
+            return
 
         dvc_exists = self.dvctree and self.dvctree.exists(top)
         repo_exists = self.repo.tree.exists(top)
         if dvc_exists and not repo_exists:
-            yield from self.dvctree.walk(top, topdown=topdown, **kwargs)
+            yield from self.dvctree.walk(
+                top, topdown=topdown, onerror=onerror, **kwargs
+            )
             return
         if repo_exists and not dvc_exists:
-            yield from self.repo.tree.walk(top, topdown=topdown)
+            yield from self.repo.tree.walk(
+                top, topdown=topdown, onerror=onerror
+            )
             return
-        if not dvc_exists and not repo_exists:
-            raise FileNotFoundError
 
         dvc_walk = self.dvctree.walk(top, topdown=topdown, **kwargs)
         repo_walk = self.repo.tree.walk(top, topdown=topdown)

dvc/scm/git/tree.py

@@ -128,17 +128,22 @@ class GitTree(BaseTree):
         if not topdown:
             yield os.path.normpath(tree.abspath), dirs, nondirs
 
-    def walk(self, top, topdown=True):
+    def walk(self, top, topdown=True, onerror=None):
         """Directory tree generator.
 
         See `os.walk` for the docs. Differences:
         - no support for symlinks
-        - it could raise exceptions, there is no onerror argument
         """
 
         tree = self.git_object_by_path(top)
         if tree is None:
-            raise OSError(errno.ENOENT, "No such file")
+            if onerror is not None:
+                onerror(OSError(errno.ENOENT, "No such file", top))
+            return
+        if tree.mode != GIT_MODE_DIR:
+            if onerror is not None:
+                onerror(NotADirectoryError(top))
+            return
 
         yield from self._walk(tree, topdown=topdown)
 

dvc/scm/tree.py

@@ -21,12 +21,11 @@ class BaseTree:
     def isfile(self, path):
         """Test whether a path is a regular file"""
 
-    def walk(self, top, topdown=True):
+    def walk(self, top, topdown=True, onerror=None):
         """Directory tree generator.
 
         See `os.walk` for the docs. Differences:
         - no support for symlinks
-        - it could raise exceptions, there is no onerror argument
         """
 
     def walk_files(self, top):
@@ -64,17 +63,12 @@ class WorkingTree(BaseTree):
         """Test whether a path is a regular file"""
         return os.path.isfile(path)
 
-    def walk(self, top, topdown=True):
+    def walk(self, top, topdown=True, onerror=None):
         """Directory tree generator.
 
         See `os.walk` for the docs. Differences:
         - no support for symlinks
-        - it could raise exceptions, there is no onerror argument
         """
-
-        def onerror(e):
-            raise e
-
         for root, dirs, files in os.walk(
             top, topdown=topdown, onerror=onerror
         ):

tests/unit/repo/test_repo_tree.py

@@ -145,6 +145,28 @@ def test_walk(tmp_dir, dvc, dvcfiles, extra_expected):
     assert len(actual) == len(expected)
 
 
+def test_walk_onerror(tmp_dir, dvc):
+    def onerror(exc):
+        raise exc
+
+    tmp_dir.dvc_gen("foo", "foo")
+    tree = RepoTree(dvc)
+
+    # path does not exist
+    for _ in tree.walk("dir"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("dir", onerror=onerror):
+            pass
+
+    # path is not a directory
+    for _ in tree.walk("foo"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("foo", onerror=onerror):
+            pass
+
+
 def test_isdvc(tmp_dir, dvc):
     tmp_dir.gen({"foo": "foo", "bar": "bar", "dir": {"baz": "baz"}})
     dvc.add("foo")

tests/unit/repo/test_tree.py

@@ -173,6 +173,28 @@ def test_walk_dir(tmp_dir, dvc, fetch, expected):
     assert len(actual) == len(expected)
 
 
+def test_walk_onerror(tmp_dir, dvc):
+    def onerror(exc):
+        raise exc
+
+    tmp_dir.dvc_gen("foo", "foo")
+    tree = DvcTree(dvc)
+
+    # path does not exist
+    for _ in tree.walk("dir"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("dir", onerror=onerror):
+            pass
+
+    # path is not a directory
+    for _ in tree.walk("foo"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("foo", onerror=onerror):
+            pass
+
+
 def test_isdvc(tmp_dir, dvc):
     tmp_dir.gen({"foo": "foo", "bar": "bar"})
     dvc.add("foo")

tests/unit/scm/test_git.py

@@ -1,5 +1,7 @@
 import os
 
+import pytest
+
 from tests.basic_env import TestDvcGit
 
 
@@ -37,6 +39,30 @@ def test_walk_with_submodules(tmp_dir, scm, git_dir):
     assert set(files) == {".gitmodules", "submodule"}
 
 
+def test_walk_onerror(tmp_dir, scm):
+    def onerror(exc):
+        raise exc
+
+    tmp_dir.scm_gen(
+        {"foo": "foo"}, commit="init",
+    )
+    tree = scm.get_tree("HEAD")
+
+    # path does not exist
+    for _ in tree.walk("dir"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("dir", onerror=onerror):
+            pass
+
+    # path is not a directory
+    for _ in tree.walk("foo"):
+        pass
+    with pytest.raises(OSError):
+        for _ in tree.walk("foo", onerror=onerror):
+            pass
+
+
 def test_is_tracked(tmp_dir, scm):
     tmp_dir.scm_gen(
         {