diff --git a/dvc/commands/experiments/save.py b/dvc/commands/experiments/save.py index c73123fea..cc75d1e50 100644 --- a/dvc/commands/experiments/save.py +++ b/dvc/commands/experiments/save.py @@ -1,6 +1,6 @@ import argparse -from dvc.cli import formatter +from dvc.cli import completion, formatter from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link from dvc.exceptions import DvcException @@ -14,6 +14,7 @@ class CmdExperimentsSave(CmdBase): def run(self): try: ref = self.repo.experiments.save( + targets=self.args.targets, name=self.args.name, force=self.args.force, include_untracked=self.args.include_untracked, @@ -41,6 +42,18 @@ def add_parser(experiments_subparsers, parent_parser): help=EXPERIMENTS_SAVE_HELP, formatter_class=formatter.RawDescriptionHelpFormatter, ) + save_parser.add_argument( + "targets", + nargs="*", + help="""\ +Stages to save. 'dvc.yaml' by default. +The targets can be path to a dvc.yaml file or `.dvc` file, +or a stage name from dvc.yaml file from +current working directory. To save a stage from dvc.yaml +from other directories, the target must be a path followed by colon `:` +and then the stage name name. +""", + ).complete = completion.DVCFILES_AND_STAGE save_parser.add_argument( "-f", "--force", diff --git a/dvc/repo/experiments/executor/base.py b/dvc/repo/experiments/executor/base.py index c168ea0b5..3a07f7fd2 100644 --- a/dvc/repo/experiments/executor/base.py +++ b/dvc/repo/experiments/executor/base.py @@ -268,6 +268,7 @@ class BaseExecutor(ABC): def save( cls, info: "ExecutorInfo", + targets: Optional[Iterable[str]] = None, force: bool = False, include_untracked: Optional[List[str]] = None, message: Optional[str] = None, @@ -293,7 +294,12 @@ class BaseExecutor(ABC): include_untracked.append(LOCK_FILE) try: - stages = dvc.commit([], force=True, relink=False) + stages = [] + if targets: + for target in targets: + stages.append(dvc.commit(target, force=True, relink=False)) + else: + stages = dvc.commit([], force=True, relink=False) exp_hash = cls.hash_exp(stages) if include_untracked: dvc.scm.add(include_untracked, force=True) # type: ignore[call-arg] diff --git a/dvc/repo/experiments/save.py b/dvc/repo/experiments/save.py index ebe5dddc0..ed067d134 100644 --- a/dvc/repo/experiments/save.py +++ b/dvc/repo/experiments/save.py @@ -1,5 +1,5 @@ import os -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Iterable, List, Optional from funcy import first @@ -14,6 +14,7 @@ logger = logger.getChild(__name__) def save( repo: "Repo", + targets: Optional[Iterable[str]] = None, name: Optional[str] = None, force: bool = False, include_untracked: Optional[List[str]] = None, @@ -32,6 +33,7 @@ def save( try: save_result = executor.save( executor.info, + targets=targets, force=force, include_untracked=include_untracked, message=message, diff --git a/tests/func/experiments/test_save.py b/tests/func/experiments/test_save.py index 52346e43e..585f0ec80 100644 --- a/tests/func/experiments/test_save.py +++ b/tests/func/experiments/test_save.py @@ -173,3 +173,19 @@ def test_exp_save_custom_message(tmp_dir, dvc, scm): exp = dvc.experiments.save(message="custom commit message") assert scm.gitpython.repo.commit(exp).message == "custom commit message" + + +def test_exp_save_target(tmp_dir, dvc, scm): + setup_stage(tmp_dir, dvc, scm) + orig_dvclock = (tmp_dir / "dvc.lock").read_text() + (tmp_dir / "bar").write_text("modified") + + tmp_dir.dvc_gen({"file": "orig"}, commit="add files") + orig_dvcfile = (tmp_dir / "file.dvc").read_text() + (tmp_dir / "file").write_text("modified") + + dvc.experiments.save(["file"]) + assert (tmp_dir / "bar").read_text() == "modified" + assert (tmp_dir / "dvc.lock").read_text() == orig_dvclock + assert (tmp_dir / "file").read_text() == "modified" + assert (tmp_dir / "file.dvc").read_text() != orig_dvcfile diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index c102e5730..0dbb00a7c 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -475,7 +475,7 @@ def test_experiments_rename_invalid(dvc, scm, mocker, capsys, caplog): def test_experiments_save(dvc, scm, mocker): - cli_args = parse_args(["exp", "save", "--name", "exp-name", "--force"]) + cli_args = parse_args(["exp", "save", "target", "--name", "exp-name", "--force"]) assert cli_args.func == CmdExperimentsSave cmd = cli_args.func(cli_args) @@ -484,7 +484,12 @@ def test_experiments_save(dvc, scm, mocker): assert cmd.run() == 0 m.assert_called_once_with( - cmd.repo, name="exp-name", force=True, include_untracked=[], message=None + cmd.repo, + targets=["target"], + name="exp-name", + force=True, + include_untracked=[], + message=None, ) @@ -500,6 +505,7 @@ def test_experiments_save_message(dvc, scm, mocker, flag): m.assert_called_once_with( cmd.repo, + targets=[], name=None, force=False, include_untracked=[],