refactoring

dvc/commands/experiments/__init__.py

@@ -62,7 +62,8 @@ def add_keep_selection_flag(experiments_subcmd_parser):
         "--keep",
         action="store_true",
         default=False,
-        help="Keep the selected experiments instead of removing them (use it with `--rev` and `--num` or with experiment names).",
+        help="Keep the selected experiments instead of removing them (use it with `--rev`'"
+             " and `--num` or with experiment names)."
     )
 def add_rev_selection_flags(
     experiments_subcmd_parser, command: str, default: bool = True

dvc/repo/experiments/remove.py

@@ -48,15 +48,7 @@ def remove(  # noqa: C901, PLR0912
     if keep_selected:
         # In keep_selected mode, identify all experiments and remove the unselected ones
         all_exp_refs = exp_refs(repo.scm, git_remote)
-
-        if exp_names:
-            selected_exp_names = set(exp_names) if isinstance(exp_names, list) else {exp_names}
-        elif rev:
-            selected_exp_names = set(
-                _resolve_exp_by_baseline(repo, [rev] if isinstance(rev, str) else rev, num, git_remote).keys()
-            )
-        else:
-            selected_exp_names = set()
+        selected_exp_names = resolve_selected_exp_names(exp_names, git_remote, num, repo, rev)
 
         # Identify experiments to remove: all experiments - selected experiments
         unselected_exp_refs = [ref for ref in all_exp_refs if ref.name not in selected_exp_names]
@@ -112,6 +104,17 @@ def remove(  # noqa: C901, PLR0912
     return removed
 
 
+def resolve_selected_exp_names(exp_names, git_remote, num, repo, rev):
+    selected_exp_names = set()
+    if exp_names:
+        selected_exp_names = set(exp_names) if isinstance(exp_names, list) else {exp_names}
+    elif rev:
+        selected_exp_names = set(
+            _resolve_exp_by_baseline(repo, [rev] if isinstance(rev, str) else rev, num, git_remote).keys()
+        )
+    return selected_exp_names
+
+
 def _resolve_exp_by_baseline(
     repo: "Repo",
     rev: list[str],

tests/func/experiments/test_remove.py

@@ -284,7 +284,6 @@ def test_keep_selected_by_nonexistent_name(tmp_dir, scm, dvc, exp_stage):
 
 def test_keep_selected_by_rev(tmp_dir, scm, dvc, exp_stage):
     # Setup: Run experiments and commit
-    baseline = scm.get_rev()
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"], name="exp1")
     exp1_ref = first(exp_refs_by_rev(scm, first(results)))
     scm.commit("commit1")
@@ -307,13 +306,10 @@ def test_keep_selected_by_rev(tmp_dir, scm, dvc, exp_stage):
 
 def test_keep_selected_by_rev_multiple(tmp_dir, scm, dvc, exp_stage):
     # Setup: Run experiments and commit
-    baseline = scm.get_rev()
-    exp1_rev = scm.get_rev()
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"], name="exp1")
     exp1_ref = first(exp_refs_by_rev(scm, first(results)))
     scm.commit("commit1")
 
-    exp2_rev = scm.get_rev()
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"], name="exp2")
     exp2_ref = first(exp_refs_by_rev(scm, first(results)))
     scm.commit("commit2")