diff --git a/dvc/commands/stage.py b/dvc/commands/stage.py index f3c11aec3c..ec5906a81f 100644 --- 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,8 +120,22 @@ 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( { @@ -128,7 +143,15 @@ def run(self): "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 index 29765ddf8e..321101934d 100644 --- 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 index 32eb9aef3d..2b090a31bd 100644 --- 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)