Merge pull request iterative#1133 from efiop/master

repro: introduce -p|--pipeline option
jeremcs · Sep 19, 2018 · 2b1ea1e · 2b1ea1e
2 parents 9dab3f2 + 34faad3
commit 2b1ea1e
Show file tree

Hide file tree

Showing 4 changed files with 75 additions and 21 deletions.
diff --git a/dvc/cli.py b/dvc/cli.py
@@ -396,6 +396,13 @@ def parse_args(argv=None):
                         default=False,
                         help='Ask for confirmation before reproducing each '
                              'stage.')
+    repro_parser.add_argument(
+                        '-p',
+                        '--pipeline',
+                        action='store_true',
+                        default=False,
+                        help='Reproduce the whole pipeline that the '
+                             'specified stage file belongs to.')
     repro_parser.set_defaults(func=CmdRepro)
 
     # Remove

diff --git a/dvc/command/repro.py b/dvc/command/repro.py
@@ -20,7 +20,8 @@ def run(self):
                                        recursive=recursive,
                                        force=self.args.force,
                                        dry=self.args.dry,
-                                       interactive=self.args.interactive)
+                                       interactive=self.args.interactive,
+                                       pipeline=self.args.pipeline)
 
                 if len(stages) == 0:
                     self.project.logger.info(CmdDataStatus.UP_TO_DATE_MSG)

diff --git a/dvc/project.py b/dvc/project.py
@@ -311,39 +311,73 @@ def reproduce(self,
                   recursive=True,
                   force=False,
                   dry=False,
-                  interactive=False):
-        import networkx as nx
-
-        stage = Stage.load(self, target)
-        G = self.graph()[1]
-        stages = nx.get_node_attributes(G, 'stage')
-        node = os.path.relpath(stage.path, self.root_dir)
+                  interactive=False,
+                  pipeline=False):
 
         if not interactive:
             config = self.config
             core = config._config[config.SECTION_CORE]
             interactive = core.get(config.SECTION_CORE_INTERACTIVE, False)
 
+        targets = []
+        if pipeline:
+            stage = Stage.load(self, target)
+            node = os.path.relpath(stage.path, self.root_dir)
+            pipelines = list(filter(lambda g: node in g.nodes(),
+                                    self.pipelines()))
+            assert len(pipelines) == 1
+            G = pipelines[0]
+            for node in G.nodes():
+                if G.in_degree(node) == 0:
+                    targets.append(os.path.join(self.root_dir, node))
+        else:
+            targets.append(target)
+
         self._files_to_git_add = []
+
+        ret = []
         with self.state:
-            if recursive:
-                ret = self._reproduce_stages(G,
-                                             stages,
-                                             node,
-                                             force,
-                                             dry,
-                                             interactive)
-            else:
-                ret = self._reproduce_stage(stages,
-                                            node,
-                                            force,
-                                            dry,
-                                            interactive)
+            for target in targets:
+                stages = self._reproduce(target,
+                                         recursive=recursive,
+                                         force=force,
+                                         dry=dry,
+                                         interactive=interactive)
+                ret.extend(stages)
 
         self._remind_to_git_add()
 
         return ret
 
+    def _reproduce(self,
+                   target,
+                   recursive=True,
+                   force=False,
+                   dry=False,
+                   interactive=False):
+        import networkx as nx
+
+        stage = Stage.load(self, target)
+        G = self.graph()[1]
+        stages = nx.get_node_attributes(G, 'stage')
+        node = os.path.relpath(stage.path, self.root_dir)
+
+        if recursive:
+            ret = self._reproduce_stages(G,
+                                         stages,
+                                         node,
+                                         force,
+                                         dry,
+                                         interactive)
+        else:
+            ret = self._reproduce_stage(stages,
+                                        node,
+                                        force,
+                                        dry,
+                                        interactive)
+
+        return ret
+
     def _reproduce_stages(self, G, stages, node, force, dry, interactive):
         import networkx as nx
 

diff --git a/tests/test_repro.py b/tests/test_repro.py
@@ -205,6 +205,18 @@ def test(self):
         self.assertEqual(len(stages), 3)
 
 
+class TestReproPipeline(TestReproChangedDeepData):
+    def test(self):
+        stages = self.dvc.reproduce(self.file1_stage,
+                                    force=True,
+                                    pipeline=True)
+        self.assertEqual(len(stages), 3)
+
+    def test_cli(self):
+        ret = main(['repro', '--pipeline', '-f', self.file1_stage])
+        self.assertEqual(ret, 0)
+
+
 class TestReproLocked(TestReproChangedData):
     def test(self):
         file2 = 'file2'