diff --git a/dvc/commands/experiments/ls.py b/dvc/commands/experiments/ls.py index 69beca90e..2e50b091c 100644 --- a/dvc/commands/experiments/ls.py +++ b/dvc/commands/experiments/ls.py @@ -3,6 +3,7 @@ import logging from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link +from dvc.ui import ui logger = logging.getLogger(__name__) @@ -16,20 +17,13 @@ class CmdExperimentsList(CmdBase): num=self.args.num, git_remote=self.args.git_remote, ) - tags = self.repo.scm.describe(exps) - remained = {baseline for baseline, tag in tags.items() if tag is None} - base = "refs/heads" - ref_heads = self.repo.scm.describe(remained, base=base) for baseline in exps: - name = baseline[:7] - if tags[baseline] or ref_heads[baseline]: - name = tags[baseline] or ref_heads[baseline][len(base) + 1 :] if not name_only: - print(f"{name}:") + ui.write(f"{baseline}:") for exp_name in exps[baseline]: indent = "" if name_only else "\t" - print(f"{indent}{exp_name}") + ui.write(f"{indent}{exp_name}") return 0 diff --git a/dvc/repo/experiments/ls.py b/dvc/repo/experiments/ls.py index 62d751117..e209f49d9 100644 --- a/dvc/repo/experiments/ls.py +++ b/dvc/repo/experiments/ls.py @@ -6,7 +6,7 @@ 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 +from .utils import exp_refs_by_baseline logger = logging.getLogger(__name__) @@ -20,19 +20,22 @@ def ls( num: int = 1, git_remote: Optional[str] = None, ): - results = defaultdict(list) - 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()) + rev_set = None + if not all_commits: + revs = iter_revs(repo.scm, [rev or "HEAD"], 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] + + tags = repo.scm.describe(ref_info_dict.keys()) + remained = {baseline for baseline, tag in tags.items() if tag is None} + base = "refs/heads" + ref_heads = repo.scm.describe(remained, base=base) + + results = defaultdict(list) + for baseline in ref_info_dict: + name = baseline[:7] + if tags[baseline] or ref_heads[baseline]: + name = tags[baseline] or ref_heads[baseline][len(base) + 1 :] + results[name] = [info.name for info in ref_info_dict[baseline]] return results diff --git a/dvc/repo/experiments/refs.py b/dvc/repo/experiments/refs.py index f563732bc..d5e594e98 100644 --- a/dvc/repo/experiments/refs.py +++ b/dvc/repo/experiments/refs.py @@ -23,9 +23,7 @@ class ExpRefInfo: namespace = EXPS_NAMESPACE - def __init__( - self, baseline_sha: Optional[str] = None, name: Optional[str] = None - ): + def __init__(self, baseline_sha: str, name: Optional[str] = None): self.baseline_sha = baseline_sha self.name: str = name if name else "" @@ -33,7 +31,7 @@ class ExpRefInfo: return "/".join(self.parts) def __repr__(self): - baseline = f"'{self.baseline_sha}'" if self.baseline_sha else "None" + baseline = f"'{self.baseline_sha}'" name = f"'{self.name}'" if self.name else "None" return f"ExpRefInfo(baseline_sha={baseline}, name={name})" @@ -41,11 +39,7 @@ class ExpRefInfo: def parts(self): return ( (self.namespace,) - + ( - (self.baseline_sha[:2], self.baseline_sha[2:]) - if self.baseline_sha - else () - ) + + ((self.baseline_sha[:2], self.baseline_sha[2:])) + ((self.name,) if self.name else ()) ) @@ -54,14 +48,13 @@ class ExpRefInfo: try: parts = ref.split("/") if ( - len(parts) < 2 - or len(parts) == 3 + len(parts) < 4 or len(parts) > 5 or "/".join(parts[:2]) != EXPS_NAMESPACE ): raise InvalidExpRefError(ref) except ValueError: raise InvalidExpRefError(ref) - baseline_sha = parts[2] + parts[3] if len(parts) >= 4 else None + baseline_sha = parts[2] + parts[3] name = parts[4] if len(parts) == 5 else None return cls(baseline_sha, name) diff --git a/dvc/repo/experiments/utils.py b/dvc/repo/experiments/utils.py index 06e15b0e5..b2422d4b0 100644 --- a/dvc/repo/experiments/utils.py +++ b/dvc/repo/experiments/utils.py @@ -79,13 +79,13 @@ def exp_refs_by_rev(scm: "Git", rev: str) -> Generator[ExpRefInfo, None, None]: def exp_refs_by_baseline( - scm: "Git", revs: Set[str], url: Optional[str] = None + scm: "Git", revs: Optional[Set[str]] = None, url: Optional[str] = None ) -> Mapping[str, List[ExpRefInfo]]: """Iterate over all experiment refs with the specified baseline.""" all_exp_refs = exp_refs(scm, url) result = defaultdict(list) for ref in all_exp_refs: - if ref.baseline_sha in revs: + if revs is None or ref.baseline_sha in revs: result[ref.baseline_sha].append(ref) return result diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 1d4329b8c..12d41aac0 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -417,7 +417,7 @@ def test_packed_args_exists(tmp_dir, scm, dvc, exp_stage, caplog): assert "Temporary DVC file" in caplog.text -def _prepare_experiments(tmp_dir, scm, dvc, exp_stage): +def test_list(tmp_dir, scm, dvc, exp_stage): baseline_a = scm.get_rev() results = dvc.experiments.run(exp_stage.addressing, params=["foo=2"]) exp_a = first(results) @@ -428,54 +428,29 @@ def _prepare_experiments(tmp_dir, scm, dvc, exp_stage): ref_info_b = first(exp_refs_by_rev(scm, exp_b)) tmp_dir.scm_gen("new", "new", commit="new") - baseline_c = scm.get_rev() results = dvc.experiments.run(exp_stage.addressing, params=["foo=4"]) exp_c = first(results) ref_info_c = first(exp_refs_by_rev(scm, exp_c)) - return baseline_a, baseline_c, ref_info_a, ref_info_b, ref_info_c - - -def test_list(tmp_dir, scm, dvc, exp_stage): - ( - baseline_a, - baseline_c, - ref_info_a, - ref_info_b, - ref_info_c, - ) = _prepare_experiments(tmp_dir, scm, dvc, exp_stage) - - assert dvc.experiments.ls() == {baseline_c: [ref_info_c.name]} + assert dvc.experiments.ls() == {"master": [ref_info_c.name]} exp_list = dvc.experiments.ls(rev=ref_info_a.baseline_sha) assert {key: set(val) for key, val in exp_list.items()} == { - baseline_a: {ref_info_a.name, ref_info_b.name} + baseline_a[:7]: {ref_info_a.name, ref_info_b.name} } 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}, + baseline_a[:7]: {ref_info_a.name, ref_info_b.name}, + "master": {ref_info_c.name}, } - -def test_list_cli(tmp_dir, scm, dvc, capsys, exp_stage): - from dvc.cli import main - - baseline_a, _, ref_info_a, ref_info_b, ref_info_c = _prepare_experiments( - tmp_dir, scm, dvc, exp_stage - ) - - # Make sure that we prioritize the current branch name scm.checkout("branch", True) - - capsys.readouterr() - assert main(["exp", "list", "-A"]) == 0 - cap = capsys.readouterr() - assert set(cap.out.split()) == set( - ["branch:", baseline_a[:7] + ":"] - + [ref_info_a.name, ref_info_b.name, ref_info_c.name] - ) + exp_list = dvc.experiments.ls(all_commits=True) + assert {key: set(val) for key, val in exp_list.items()} == { + baseline_a[:7]: {ref_info_a.name, ref_info_b.name}, + "branch": {ref_info_c.name}, + } def test_subdir(tmp_dir, scm, dvc, workspace): diff --git a/tests/func/experiments/test_remote.py b/tests/func/experiments/test_remote.py index aa67d8237..93b82a14f 100644 --- a/tests/func/experiments/test_remote.py +++ b/tests/func/experiments/test_remote.py @@ -151,7 +151,6 @@ def test_list_remote(tmp_dir, scm, dvc, git_downstream, exp_stage, use_url): ref_info_b = first(exp_refs_by_rev(scm, exp_b)) tmp_dir.scm_gen("new", "new", commit="new") - baseline_c = scm.get_rev() results = dvc.experiments.run(exp_stage.addressing, params=["foo=4"]) exp_c = first(results) ref_info_c = first(exp_refs_by_rev(scm, exp_c)) @@ -165,13 +164,13 @@ def test_list_remote(tmp_dir, scm, dvc, git_downstream, exp_stage, use_url): 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} + baseline_a[:7]: {ref_info_a.name, ref_info_b.name} } 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}, + baseline_a[:7]: {ref_info_a.name, ref_info_b.name}, + "master": {ref_info_c.name}, }