feat(gc): allow num with all-branches (#11013)

dvc/commands/gc.py

@@ -40,8 +40,15 @@ class CmdGC(CmdBase):
             msg += " and all git commits"
         else:
             if self.args.all_branches and self.args.all_tags:
+                if self.args.num:
+                    msg += f" and last {self.args.num} commits from all git branches"
+                    msg += " and all git tags"
+                else:
                     msg += " and all git branches and tags"
             elif self.args.all_branches:
+                if self.args.num:
+                    msg += f" and last {self.args.num} commits from all git branches"
+                else:
                     msg += " and all git branches"
             elif self.args.all_tags:
                 msg += " and all git tags"
@@ -126,8 +133,8 @@ def add_parser(subparsers, parent_parser):
         metavar="<num>",
         help=(
             "Keep data files used in the last `num` commits "
-            "starting from the `--rev` <commit>. "
-            "Only used if `--rev` is also provided. "
+            "starting from each selected revision root. "
+            "Can be used with `--rev` and `--all-branches`. "
             "Defaults to `1`."
         ),
     )

dvc/repo/gc.py

@@ -28,8 +28,10 @@ def _validate_args(**kwargs):
             "`--all-experiments`, `--all-commits`, `--date` or `--rev` "
             "needs to be set."
         )
-    if kwargs.get("num") and not kwargs.get("rev"):
-        raise InvalidArgumentError("`--num` can only be used alongside `--rev`")
+    if kwargs.get("num") and not (kwargs.get("rev") or kwargs.get("all_branches")):
+        raise InvalidArgumentError(
+            "`--num` can only be used alongside `--rev` or `--all-branches`"
+        )
 
 
 def _used_obj_ids_not_in_remote(

dvc/scm.py

@@ -234,7 +234,7 @@ def iter_revs(
         results.extend(scm.list_all_commits())
     else:
         if all_branches:
-            results.extend(scm.list_branches())
+            results.extend(_get_n_commits(scm, scm.list_branches(), num))
 
         if all_tags:
             results.extend(scm.list_tags())
@@ -258,7 +258,10 @@ def iter_revs(
         results.extend(exp_commits(scm))
 
     rev_resolver = partial(resolve_rev, scm)
-    return group_by(rev_resolver, results)
+    grouped = group_by(rev_resolver, results)
+    for rev, names in grouped.items():
+        grouped[rev] = list(dict.fromkeys(names))
+    return grouped
 
 
 def lfs_prefetch(fs: "FileSystem", paths: list[str]):

tests/func/test_gc.py

@@ -92,6 +92,27 @@ def test_gc_multiple_dvc_repos(tmp_dir, scm, dvc, erepo_dir):
     assert len(list(odb.all())) == 2
 
 
+def test_gc_multiple_dvc_repos_all_branches_num(tmp_dir, scm, dvc, erepo_dir):
+    tmp_dir.dvc_gen("main", "main", commit="main")
+
+    erepo_dir.dvc.cache.local.path = dvc.cache.local.path
+    with erepo_dir.chdir():
+        erepo_dir.scm_gen("base", "base", commit="base")
+        with erepo_dir.branch("feature", new=True):
+            erepo_dir.dvc_gen("shared", "feature-v1", commit="feature-v1")
+            erepo_dir.dvc.remove("shared.dvc")
+            erepo_dir.dvc_gen("shared", "feature-v2", commit="feature-v2")
+
+    odb = dvc.cache.local
+    assert len(list(odb.all())) == 3
+
+    dvc.gc(all_branches=True, num=2, repos=[erepo_dir])
+    assert len(list(odb.all())) == 3
+
+    dvc.gc(all_branches=True, repos=[erepo_dir])
+    assert len(list(odb.all())) == 2
+
+
 def test_all_commits(tmp_dir, scm, dvc):
     tmp_dir.dvc_gen("testfile", "uncommitted")
     tmp_dir.dvc_gen("testfile", "committed", commit="committed")

tests/unit/command/test_gc.py

@@ -65,6 +65,30 @@ def test_(dvc, scm, mocker):
     with pytest.raises(InvalidArgumentError):
         cmd.run()
 
+    cli_args = parse_args(["gc", "--all-branches", "--num", "2", "--force"])
+    cmd = cli_args.func(cli_args)
+
+    assert cmd.run() == 0
+
+    m.assert_called_with(
+        workspace=False,
+        all_tags=False,
+        all_branches=True,
+        all_commits=False,
+        all_experiments=False,
+        commit_date=None,
+        cloud=False,
+        remote=None,
+        force=True,
+        jobs=None,
+        repos=None,
+        rev=None,
+        num=2,
+        not_in_remote=False,
+        dry=False,
+        skip_failed=False,
+    )
+
     cli_args = parse_args(["gc", "--cloud", "--not-in-remote"])
     cmd = cli_args.func(cli_args)
     with pytest.raises(InvalidArgumentError):

tests/unit/scm/test_scm.py

@@ -47,6 +47,13 @@ def test_iter_revs(tmp_dir, scm, mocker):
     assert gen == {rev_old: ["tag"]}
     gen = iter_revs(scm, all_branches=True)
     assert gen == {rev_old: [old], rev_new: ["new"], rev_other: ["other"]}
+    gen = iter_revs(scm, all_branches=True, num=2)
+    assert gen == {
+        rev_old: [old, rev_old],
+        rev_root: [rev_root],
+        rev_new: ["new"],
+        rev_other: ["other"],
+    }
     gen = iter_revs(scm, all_tags=True)
     assert gen == {rev_old: ["tag"]}
     gen = iter_revs(scm, all_commits=True)