dvc: repo: unify collect() method

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/repo/__init__.py

@@ -143,10 +143,13 @@ class Repo(object):
         assert len(pipelines) == 1
         return pipelines[0]
 
-    def collect(self, target, with_deps=False):
+    def collect(self, target, with_deps=False, recursive=False):
         import networkx as nx
         from dvc.stage import Stage
 
+        if not target or recursive:
+            return self.active_stages(target)
+
         stage = Stage.load(self, target)
         if not with_deps:
             return [stage]

dvc/repo/checkout.py

@@ -29,15 +29,14 @@ def checkout(self, target=None, with_deps=False, force=False, recursive=False):
 
     all_stages = self.stages()
 
-    if target and not recursive:
-        try:
-            stages = self.collect(target, with_deps=with_deps)
-        except (StageFileDoesNotExistError, StageFileBadNameError) as exc:
-            raise DvcException(
-                str(exc) + " Did you mean 'git checkout {}'?".format(target)
-            )
-    else:
-        stages = self.active_stages(target)
+    try:
+        stages = self.collect(target, with_deps=with_deps, recursive=recursive)
+    except (StageFileDoesNotExistError, StageFileBadNameError) as exc:
+        if not target:
+            raise
+        raise DvcException(
+            str(exc) + " Did you mean 'git checkout {}'?".format(target)
+        )
 
     with self.state:
         _cleanup_unused_links(self, all_stages)

dvc/repo/commit.py

@@ -1,9 +1,5 @@
 def commit(self, target, with_deps=False, recursive=False, force=False):
-    if target and not recursive:
-        stages = self.collect(target, with_deps=with_deps)
-    else:
-        stages = self.active_stages(target)
-
+    stages = self.collect(target, with_deps=with_deps, recursive=recursive)
     with self.state:
         for stage in stages:
             stage.check_can_commit(force=force)

dvc/repo/status.py

@@ -6,10 +6,7 @@ import dvc.logger as logger
 def _local_status(self, target=None, with_deps=False):
     status = {}
 
-    if target:
-        stages = self.collect(target, with_deps=with_deps)
-    else:
-        stages = self.active_stages()
+    stages = self.collect(target, with_deps=with_deps)
 
     for stage in stages:
         if stage.locked: