remote local, repo: use walk_files instead of os.walk (#1750)

dvc/remote/local.py

@@ -20,7 +20,15 @@ from dvc.remote.base import (
     STATUS_DELETED,
     STATUS_MISSING,
 )
-from dvc.utils import remove, move, copyfile, dict_md5, to_chunks, tmp_fname
+from dvc.utils import (
+    remove,
+    move,
+    copyfile,
+    dict_md5,
+    to_chunks,
+    tmp_fname,
+    walk_files,
+)
 from dvc.utils import LARGE_DIR_SIZE
 from dvc.config import Config
 from dvc.exceptions import DvcException
@@ -353,11 +361,7 @@ class RemoteLOCAL(RemoteBase):
         return not self.changed_cache(current_md5)
 
     def _discard_working_directory_changes(self, path, dir_info, force=False):
-        working_dir_files = set(
-            os.path.join(root, file)
-            for root, _, files in os.walk(str(path))
-            for file in files
-        )
+        working_dir_files = set(path for path in walk_files(path))
 
         cached_files = set(
             os.path.join(path, file["relpath"]) for file in dir_info

dvc/repo/unprotect.py

@@ -6,7 +6,7 @@ import uuid
 
 import dvc.logger as logger
 from dvc.system import System
-from dvc.utils import copyfile, remove
+from dvc.utils import copyfile, remove, walk_files
 from dvc.exceptions import DvcException
 
 
@@ -33,10 +33,8 @@ def _unprotect_file(path):
 
 
 def _unprotect_dir(path):
-    for root, dirs, files in os.walk(str(path)):
-        for f in files:
-            path = os.path.join(root, f)
-            _unprotect_file(path)
+    for path in walk_files(path):
+        _unprotect_file(path)
 
 
 def unprotect(path):

dvc/utils/__init__.py

@@ -242,6 +242,6 @@ def load_stage_file(path):
 
 
 def walk_files(directory):
-    for root, _, files in os.walk(directory):
+    for root, _, files in os.walk(str(directory)):
         for f in files:
             yield os.path.join(root, f)

tests/test_checkout.py

@@ -12,6 +12,7 @@ from dvc.main import main
 from dvc import progress
 from dvc.repo import Repo as DvcRepo
 from dvc.system import System
+from dvc.utils import walk_files
 from tests.basic_env import TestDvc
 from tests.test_repro import TestRepro
 from dvc.stage import Stage
@@ -126,10 +127,7 @@ class CheckoutBase(TestDvc):
         FileInfo = collections.namedtuple("FileInfo", "path inode")
 
         paths = [
-            os.path.join(root, file)
-            for output in stage.outs
-            for root, _, files in os.walk(output.path)
-            for file in files
+            path for output in stage.outs for path in walk_files(output.path)
         ]
 
         return [