diff --git a/dvc/commands/experiments/branch.py b/dvc/commands/experiments/branch.py index 6d8137513..eb06a84f8 100644 --- a/dvc/commands/experiments/branch.py +++ b/dvc/commands/experiments/branch.py @@ -26,5 +26,13 @@ def add_parser(experiments_subparsers, parent_parser): experiments_branch_parser.add_argument( "experiment", help="Experiment to be promoted." ) - experiments_branch_parser.add_argument("branch", help="Name for the new Git branch") + experiments_branch_parser.add_argument( + "branch", + nargs="?", + default=None, + help=( + "Optional name for the new Git branch. " + "Defaults to '{experiment-name}-branch'." + ), + ) experiments_branch_parser.set_defaults(func=CmdExperimentsBranch) diff --git a/dvc/repo/experiments/branch.py b/dvc/repo/experiments/branch.py index 10b946e41..32c2a52d9 100644 --- a/dvc/repo/experiments/branch.py +++ b/dvc/repo/experiments/branch.py @@ -13,7 +13,7 @@ logger = logging.getLogger(__name__) @locked @scm_context -def branch(repo, exp_rev, branch_name, *args, **kwargs): +def branch(repo, exp_rev, branch_name=None, **kwargs): from dvc.scm import resolve_rev try: @@ -46,6 +46,8 @@ def branch(repo, exp_rev, branch_name, *args, **kwargs): if not ref_info: raise InvalidExpRevError(exp_rev) + branch_name = branch_name or f"{ref_info.name}-branch" + branch_ref = f"refs/heads/{branch_name}" if repo.scm.get_ref(branch_ref): raise InvalidArgumentError(f"Git branch '{branch_name}' already exists.") diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index a96f0c332..2cb499269 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -246,11 +246,12 @@ def test_branch(tmp_dir, scm, dvc, exp_stage): with pytest.raises(InvalidArgumentError): dvc.experiments.branch("foo", "branch-exists") + dvc.experiments.branch("foo") dvc.experiments.branch("foo", "branch-name") dvc.experiments.branch(exp_a, "branch-rev") dvc.experiments.branch(ref_a, "branch-ref") - for name in ["branch-name", "branch-rev", "branch-ref"]: + for name in ["foo-branch", "branch-name", "branch-rev", "branch-ref"]: assert name in scm.list_branches() assert scm.resolve_rev(name) == exp_a diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 8fdbb3109..9ca3fd54f 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -187,15 +187,21 @@ def test_experiments_gc(dvc, scm, mocker): def test_experiments_branch(dvc, scm, mocker): - cli_args = parse_args(["experiments", "branch", "expname", "branchname"]) + m = mocker.patch("dvc.repo.experiments.branch.branch", return_value={}) + + cli_args = parse_args(["experiments", "branch", "expname"]) assert cli_args.func == CmdExperimentsBranch cmd = cli_args.func(cli_args) - m = mocker.patch("dvc.repo.experiments.branch.branch", return_value={}) - assert cmd.run() == 0 - m.assert_called_once_with(cmd.repo, "expname", "branchname") + m.assert_called_with(cmd.repo, "expname", None) + + cli_args = parse_args(["experiments", "branch", "expname", "branchname"]) + cmd = cli_args.func(cli_args) + assert cmd.run() == 0 + + m.assert_called_with(cmd.repo, "expname", "branchname") def test_experiments_list(dvc, scm, mocker):