exp list: add the new flags (#7153)

fix: #7153
Unify the flags in `exp list` to `exp show`

1. Add new `-A`, `-n` flags to `exp list` command.
2. Add unit test for the flag parse.
3. unify the revs retrive logic to `exp show`.
4. Add functional test for `exp list`
5. Update the documents.

Co-authored-by: Jorge Orpinel <[email protected]>

dvc/commands/experiments/__init__.py

@@ -58,3 +58,43 @@ def add_parser(subparsers, parent_parser):
     for cmd in SUB_COMMANDS:
         cmd.add_parser(experiments_subparsers, parent_parser)
     fix_plumbing_subparsers(experiments_subparsers)
+
+
+def add_rev_selection_flags(
+    experiments_subcmd_parser, command: str, default: bool = True
+):
+    experiments_subcmd_parser.add_argument(
+        "-A",
+        "--all-commits",
+        action="store_true",
+        default=False,
+        help=(
+            f"{command} all experiments in the repository "
+            "(overrides `--rev` and `--num`)."
+        ),
+    )
+    default_msg = " (HEAD by default)" if default else ""
+    msg = (
+        f"{command} experiments derived from the specified `<commit>` as "
+        f"baseline{default_msg}."
+    )
+    experiments_subcmd_parser.add_argument(
+        "--rev",
+        type=str,
+        default=None,
+        help=msg,
+        metavar="<commit>",
+    )
+    experiments_subcmd_parser.add_argument(
+        "-n",
+        "--num",
+        type=int,
+        default=1,
+        dest="num",
+        metavar="<num>",
+        help=(
+            f"{command} experiments from the `--rev` baseline and from `num` "
+            "commits before it (first parents). Give a negative value to "
+            "include all first-parent commits."
+        ),
+    )

dvc/commands/experiments/ls.py

@@ -11,9 +11,10 @@ class CmdExperimentsList(CmdBase):
     def run(self):
         names_only = self.args.names_only
         exps = self.repo.experiments.ls(
+            all_commits=self.args.all_commits,
             rev=self.args.rev,
+            num=self.args.num,
             git_remote=self.args.git_remote,
-            all_=self.args.all,
         )
         for baseline in exps:
             tag = self.repo.scm.describe(baseline)
@@ -32,6 +33,7 @@ class CmdExperimentsList(CmdBase):
 
 
 def add_parser(experiments_subparsers, parent_parser):
+    from . import add_rev_selection_flags
 
     EXPERIMENTS_LIST_HELP = "List local and remote experiments."
     experiments_list_parser = experiments_subparsers.add_parser(
@@ -41,19 +43,7 @@ def add_parser(experiments_subparsers, parent_parser):
         help=EXPERIMENTS_LIST_HELP,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    experiments_list_parser.add_argument(
-        "--rev",
-        type=str,
-        default=None,
-        help=(
-            "List experiments derived from the specified revision. "
-            "Defaults to HEAD if neither `--rev` nor `--all` are specified."
-        ),
-        metavar="<rev>",
-    )
-    experiments_list_parser.add_argument(
-        "--all", action="store_true", help="List all experiments."
-    )
+    add_rev_selection_flags(experiments_list_parser, "List")
     experiments_list_parser.add_argument(
         "--names-only",
         action="store_true",

dvc/commands/experiments/show.py

@@ -515,6 +515,8 @@ class CmdExperimentsShow(CmdBase):
 
 
 def add_parser(experiments_subparsers, parent_parser):
+    from . import add_rev_selection_flags
+
     EXPERIMENTS_SHOW_HELP = "Print experiments."
     experiments_show_parser = experiments_subparsers.add_parser(
         "show",
@@ -523,6 +525,7 @@ def add_parser(experiments_subparsers, parent_parser):
         help=EXPERIMENTS_SHOW_HELP,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
+    add_rev_selection_flags(experiments_show_parser, "Show")
     experiments_show_parser.add_argument(
         "-a",
         "--all-branches",
@@ -537,32 +540,6 @@ def add_parser(experiments_subparsers, parent_parser):
         default=False,
         help="Show experiments derived from all Git tags.",
     )
-    experiments_show_parser.add_argument(
-        "-A",
-        "--all-commits",
-        action="store_true",
-        default=False,
-        help="Show experiments derived from all Git commits.",
-    )
-    experiments_show_parser.add_argument(
-        "--rev",
-        type=str,
-        default=None,
-        help=(
-            "Show experiments derived from the specified revision. "
-            "Defaults to HEAD if none of `--rev`,`-a`,`-A`,`-T` is specified."
-        ),
-        metavar="<rev>",
-    )
-    experiments_show_parser.add_argument(
-        "-n",
-        "--num",
-        type=int,
-        default=1,
-        dest="num",
-        metavar="<num>",
-        help="Show the last `num` commits from <rev>.",
-    )
     experiments_show_parser.add_argument(
         "--no-pager",
         action="store_true",

dvc/repo/experiments/ls.py

@@ -3,49 +3,36 @@ from collections import defaultdict
 
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
+from dvc.scm import iter_revs
+from dvc.types import Optional
 
-from .utils import (
-    exp_refs,
-    exp_refs_by_baseline,
-    remote_exp_refs,
-    remote_exp_refs_by_baseline,
-)
+from .utils import exp_refs, exp_refs_by_baseline
 
 logger = logging.getLogger(__name__)
 
 
 @locked
 @scm_context
-def ls(repo, *args, rev=None, git_remote=None, all_=False, **kwargs):
-    from scmrepo.git import Git
-
-    from dvc.scm import RevError, resolve_rev
-
-    if rev:
-        try:
-            rev = resolve_rev(repo.scm, rev)
-        except RevError:
-            if not (git_remote and Git.is_sha(rev)):
-                # This could be a remote rev that has not been fetched yet
-                raise
-    elif not all_:
-        rev = repo.scm.get_rev()
-
+def ls(
+    repo,
+    rev: Optional[str] = None,
+    all_commits: bool = False,
+    num: int = 1,
+    git_remote: Optional[str] = None,
+):
     results = defaultdict(list)
-
-    if rev:
-        if git_remote:
-            gen = remote_exp_refs_by_baseline(repo.scm, git_remote, rev)
-        else:
-            gen = exp_refs_by_baseline(repo.scm, rev)
-        for info in gen:
-            results[rev].append(info.name)
-    elif all_:
-        if git_remote:
-            gen = remote_exp_refs(repo.scm, git_remote)
-        else:
-            gen = exp_refs(repo.scm)
+    if all_commits:
+        gen = exp_refs(repo.scm, git_remote)
         for info in gen:
             results[info.baseline_sha].append(info.name)
+        return results
+
+    rev = rev or "HEAD"
+
+    revs = iter_revs(repo.scm, [rev], num)
+    rev_set = set(revs.keys())
+    ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
+    for rev, ref_info_list in ref_info_dict.items():
+        results[rev] = [ref_info.name for ref_info in ref_info_list]
 
     return results

dvc/repo/experiments/utils.py

@@ -5,6 +5,7 @@ from typing import (
     Generator,
     Iterable,
     List,
+    Mapping,
     Optional,
     Set,
     Union,
@@ -24,9 +25,16 @@ from .base import (
 )
 
 
-def exp_refs(scm: "Git") -> Generator["ExpRefInfo", None, None]:
+def exp_refs(
+    scm: "Git", url: Optional[str] = None
+) -> Generator["ExpRefInfo", None, None]:
     """Iterate over all experiment refs."""
-    for ref in scm.iter_refs(base=EXPS_NAMESPACE):
+    ref_gen = (
+        iter_remote_refs(scm, url, base=EXPS_NAMESPACE)
+        if url
+        else scm.iter_refs(base=EXPS_NAMESPACE)
+    )
+    for ref in ref_gen:
         if ref.startswith(EXEC_NAMESPACE) or ref == EXPS_STASH:
             continue
         yield ExpRefInfo.from_ref(ref)
@@ -42,14 +50,15 @@ def exp_refs_by_rev(
 
 
 def exp_refs_by_baseline(
-    scm: "Git", rev: str
-) -> Generator["ExpRefInfo", None, None]:
+    scm: "Git", revs: Set[str], url: Optional[str] = None
+) -> Mapping[str, List[ExpRefInfo]]:
     """Iterate over all experiment refs with the specified baseline."""
-    ref_info = ExpRefInfo(baseline_sha=rev)
-    for ref in scm.iter_refs(base=str(ref_info)):
-        if ref.startswith(EXEC_NAMESPACE) or ref == EXPS_STASH:
-            continue
-        yield ExpRefInfo.from_ref(ref)
+    all_exp_refs = exp_refs(scm, url)
+    result = defaultdict(list)
+    for ref in all_exp_refs:
+        if ref.baseline_sha in revs:
+            result[ref.baseline_sha].append(ref)
+    return result
 
 
 def iter_remote_refs(
@@ -88,22 +97,12 @@ def push_refspec(
         raise GitAuthError(str(exc))
 
 
-def remote_exp_refs(
-    scm: "Git", url: str
-) -> Generator["ExpRefInfo", None, None]:
-    """Iterate over all remote experiment refs."""
-    for ref in iter_remote_refs(scm, url, base=EXPS_NAMESPACE):
-        if ref.startswith(EXEC_NAMESPACE) or ref == EXPS_STASH:
-            continue
-        yield ExpRefInfo.from_ref(ref)
-
-
 def exp_refs_by_names(
     scm: "Git", names: Set[str], url: Optional[str] = None
 ) -> Dict[str, List[ExpRefInfo]]:
     """Iterate over all experiment refs matching the specified names."""
     resolve_results = defaultdict(list)
-    ref_info_gen = remote_exp_refs(scm, url) if url else exp_refs(scm)
+    ref_info_gen = exp_refs(scm, url)
     for ref_info in ref_info_gen:
         if ref_info.name in names:
             resolve_results[ref_info.name].append(ref_info)
@@ -111,17 +110,6 @@ def exp_refs_by_names(
     return resolve_results
 
 
-def remote_exp_refs_by_baseline(
-    scm: "Git", url: str, rev: str
-) -> Generator["ExpRefInfo", None, None]:
-    """Iterate over all remote experiment refs with the specified baseline."""
-    ref_info = ExpRefInfo(baseline_sha=rev)
-    for ref in iter_remote_refs(scm, url, base=str(ref_info)):
-        if ref.startswith(EXEC_NAMESPACE) or ref == EXPS_STASH:
-            continue
-        yield ExpRefInfo.from_ref(ref)
-
-
 def exp_commits(
     scm: "Git", ref_infos: Iterable["ExpRefInfo"] = None
 ) -> Iterable[str]:

dvc/scm.py

@@ -8,7 +8,7 @@ from scmrepo.base import Base  # noqa: F401, pylint: disable=unused-import
 from scmrepo.git import Git
 from scmrepo.noscm import NoSCM
 
-from dvc.exceptions import DvcException, InvalidArgumentError
+from dvc.exceptions import DvcException
 from dvc.progress import Tqdm
 
 if TYPE_CHECKING:
@@ -136,7 +136,7 @@ def resolve_rev(scm: "Git", rev: str) -> str:
 
 def iter_revs(
     scm: "Git",
-    head_revs: Optional[List[str]] = None,
+    revs: Optional[List[str]] = None,
     num: int = 1,
     all_branches: bool = False,
     all_tags: bool = False,
@@ -144,42 +144,39 @@ def iter_revs(
     all_experiments: bool = False,
 ) -> Mapping[str, List[str]]:
 
-    if num < 1 and num != -1:
-        raise InvalidArgumentError(f"Invalid number of commits '{num}'")
-
-    if not any(
-        [head_revs, all_branches, all_tags, all_commits, all_experiments]
-    ):
+    if not any([revs, all_branches, all_tags, all_commits, all_experiments]):
         return {}
 
-    head_revs = head_revs or []
-    revs = []
-    for rev in head_revs:
-        revs.append(rev)
+    revs = revs or []
+    results = []
+    for rev in revs:
+        if num == 0:
+            continue
+        results.append(rev)
         n = 1
         while True:
             if num == n:
                 break
             try:
                 head = f"{rev}~{n}"
-                revs.append(resolve_rev(scm, head))
+                results.append(resolve_rev(scm, head))
             except RevError:
                 break
             n += 1
 
     if all_commits:
-        revs.extend(scm.list_all_commits())
+        results.extend(scm.list_all_commits())
     else:
         if all_branches:
-            revs.extend(scm.list_branches())
+            results.extend(scm.list_branches())
 
         if all_tags:
-            revs.extend(scm.list_tags())
+            results.extend(scm.list_tags())
 
     if all_experiments:
         from dvc.repo.experiments.utils import exp_commits
 
-        revs.extend(exp_commits(scm))
+        results.extend(exp_commits(scm))
 
     rev_resolver = partial(resolve_rev, scm)
-    return group_by(rev_resolver, revs)
+    return group_by(rev_resolver, results)

tests/func/experiments/test_experiments.py

@@ -479,7 +479,7 @@ def test_list(tmp_dir, scm, dvc, exp_stage):
         baseline_a: {ref_info_a.name, ref_info_b.name}
     }
 
-    exp_list = dvc.experiments.ls(all_=True)
+    exp_list = dvc.experiments.ls(all_commits=True)
     assert {key: set(val) for key, val in exp_list.items()} == {
         baseline_a: {ref_info_a.name, ref_info_b.name},
         baseline_c: {ref_info_c.name},

tests/func/experiments/test_remote.py

@@ -132,12 +132,13 @@ def test_list_remote(tmp_dir, scm, dvc, git_downstream, exp_stage, use_url):
     downstream_exp = git_downstream.tmp_dir.dvc.experiments
     assert downstream_exp.ls(git_remote=remote) == {}
 
+    git_downstream.tmp_dir.scm.fetch_refspecs(remote, ["master:master"])
     exp_list = downstream_exp.ls(rev=baseline_a, git_remote=remote)
     assert {key: set(val) for key, val in exp_list.items()} == {
         baseline_a: {ref_info_a.name, ref_info_b.name}
     }
 
-    exp_list = downstream_exp.ls(all_=True, git_remote=remote)
+    exp_list = downstream_exp.ls(all_commits=True, git_remote=remote)
     assert {key: set(val) for key, val in exp_list.items()} == {
         baseline_a: {ref_info_a.name, ref_info_b.name},
         baseline_c: {ref_info_c.name},

tests/func/experiments/test_show.py

@@ -7,7 +7,6 @@ import pytest
 from funcy import first, get_in
 
 from dvc.cli import main
-from dvc.exceptions import InvalidArgumentError
 from dvc.repo.experiments.base import EXPS_STASH, ExpRefInfo
 from dvc.repo.experiments.executor.base import (
     EXEC_PID_DIR,
@@ -287,9 +286,7 @@ def test_show_multiple_commits(tmp_dir, scm, dvc, exp_stage):
     tmp_dir.scm_gen("file", "file", "commit")
     next_rev = scm.get_rev()
 
-    dvc.experiments.show(num=-1)
-    with pytest.raises(InvalidArgumentError):
-        dvc.experiments.show(num=-2)
+    dvc.experiments.show(num=-2)
 
     expected = {"workspace", init_rev, next_rev}
     results = dvc.experiments.show(num=2)

tests/unit/command/test_experiments.py

@@ -197,9 +197,11 @@ def test_experiments_list(dvc, scm, mocker):
             "experiments",
             "list",
             "origin",
+            "--all-commits",
+            "-n",
+            "-1",
             "--rev",
             "foo",
-            "--all",
             "--names-only",
         ]
     )
@@ -211,7 +213,11 @@ def test_experiments_list(dvc, scm, mocker):
     assert cmd.run() == 0
 
     m.assert_called_once_with(
-        cmd.repo, git_remote="origin", rev="foo", all_=True
+        cmd.repo,
+        git_remote="origin",
+        rev="foo",
+        all_commits=True,
+        num=-1,
     )