Merge pull request #3048 from Suor/optimize-stages

dvc: optimize and clean up stages collection
iterative · Jan 3, 2020 · 85c35f8 · 85c35f8
2 parents bdcc995 + b2d083f
commit 85c35f8
Showing 1 changed file with 5 additions and 15 deletions.
diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py
@@ -363,17 +363,6 @@ def graph(self):
     def pipelines(self):
         return get_pipelines(self.graph)
 
-    @staticmethod
-    def _filter_out_dirs(dirs, outs, root_dir):
-        def filter_dirs(dname):
-            path = os.path.join(root_dir, dname)
-            for out in outs:
-                if path == os.path.normpath(out):
-                    return False
-            return True
-
-        return list(filter(filter_dirs, dirs))
-
     @cached_property
     def stages(self):
         """
@@ -388,20 +377,21 @@ def stages(self):
         from dvc.stage import Stage
 
         stages = []
-        outs = []
+        outs = set()
 
         for root, dirs, files in self.tree.walk(self.root_dir):
             for fname in files:
                 path = os.path.join(root, fname)
                 if not Stage.is_valid_filename(path):
                     continue
                 stage = Stage.load(self, path)
+                stages.append(stage)
+
                 for out in stage.outs:
                     if out.scheme == "local":
-                        outs.append(out.fspath + out.sep)
-                stages.append(stage)
+                        outs.add(out.fspath)
 
-            dirs[:] = self._filter_out_dirs(dirs, outs, root)
+            dirs[:] = [d for d in dirs if os.path.join(root, d) not in outs]
 
         return stages