stage add: introduce --run option (#9480)

* stage add: introduce --run option Part of #5846. * fix ui message
iterative · May 22, 2023 · 1989c0f · 1989c0f
1 parent 3ddd4b8
commit 1989c0f
Show file tree

Hide file tree

Showing 3 changed files with 51 additions and 11 deletions.
diff --git a/dvc/commands/stage.py b/dvc/commands/stage.py
@@ -1,5 +1,6 @@
 import argparse
 import logging
+from contextlib import contextmanager
 from itertools import chain, filterfalse
 from typing import TYPE_CHECKING, Dict, Iterable, List
 
@@ -119,16 +120,38 @@ def quote_argument(arg: str):
     return " ".join(map(quote_argument, commands))
 
 
+@contextmanager
+def _disable_logging(highest_level=logging.CRITICAL):
+    previous_level = logging.root.manager.disable
+
+    logging.disable(highest_level)
+
+    try:
+        yield
+    finally:
+        logging.disable(previous_level)
+
+
 class CmdStageAdd(CmdBase):
     def run(self):
+        from dvc.repo import lock_repo
+
         kwargs = vars(self.args)
         kwargs.update(
             {
                 "cmd": parse_cmd(kwargs.pop("command")),
                 "params": parse_params(self.args.params),
             }
         )
-        self.repo.stage.add(**kwargs)
+
+        with self.repo.scm_context, lock_repo(self.repo):
+            with _disable_logging(logging.INFO):
+                stage = self.repo.stage.add(**kwargs)
+            logger.info("Added stage %r in %r", stage.addressing, stage.relpath)
+            if self.args.run:
+                stage.run()
+                stage.dump(update_pipeline=False)
+
         return 0
 
 
@@ -257,6 +280,12 @@ def _add_common_args(parser):
             "This doesn't affect any DVC operations."
         ),
     )
+    parser.add_argument(
+        "--run",
+        action="store_true",
+        default=False,
+        help="Execute the stage after generating it.",
+    )
     parser.add_argument(
         "command",
         nargs=argparse.REMAINDER,

diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py
@@ -122,16 +122,15 @@ def add(
             force=force,
             **stage_data,
         )
-        with self.repo.scm_context:
-            stage.dump(update_lock=update_lock)
-            try:
-                stage.ignore_outs()
-            except FileNotFoundError as exc:
-                ui.warn(
-                    f"Could not create .gitignore entry in {exc.filename}."
-                    " DVC will attempt to create .gitignore entry again when"
-                    " the stage is run."
-                )
+        stage.dump(update_lock=update_lock)
+        try:
+            stage.ignore_outs()
+        except FileNotFoundError as exc:
+            ui.warn(
+                f"Could not create .gitignore entry in {exc.filename}."
+                " DVC will attempt to create .gitignore entry again when"
+                " the stage is run."
+            )
 
         return stage
 

diff --git a/tests/unit/command/test_stage.py b/tests/unit/command/test_stage.py
@@ -86,3 +86,15 @@ def test_stage_add(mocker, dvc, command, parsed_command):
         cmd=parsed_command,
         force=True,
     )
+
+
+def test_stage_add_and_run(mocker, dvc):
+    cli_args = parse_args(["stage", "add", "--run", "-n", "foo", "-o", "foo", "cmd"])
+    cmd = cli_args.func(cli_args)
+    add_mock = mocker.patch.object(cmd.repo.stage, "add")
+
+    assert cmd.run() == 0
+
+    assert called_once_with_subset(add_mock, name="foo", outs=["foo"], cmd="cmd")
+    add_mock.return_value.run.assert_called_once()
+    add_mock.return_value.dump.assert_called_once_with(update_pipeline=False)