Add 'pack-refs' command

Generally speaking, packed refs are much faster to read than
file-based refs, when present in large numbers.

dulwich/cli.py

@@ -300,6 +300,18 @@ class cmd_symbolic_ref(Command):
         porcelain.symbolic_ref(".", ref_name=ref_name, force="--force" in args)
 
 
+class cmd_pack_refs(Command):
+    def run(self, argv):
+        parser = argparse.ArgumentParser()
+        parser.add_argument('--all', action='store_true')
+        # ignored, we never prune
+        parser.add_argument('--no-prune', action='store_true')
+
+        args = parser.parse_args(argv)
+
+        porcelain.pack_refs(".", all=args.all)
+
+
 class cmd_show(Command):
     def run(self, argv):
         parser = argparse.ArgumentParser()
@@ -756,6 +768,7 @@ commands = {
     "ls-remote": cmd_ls_remote,
     "ls-tree": cmd_ls_tree,
     "pack-objects": cmd_pack_objects,
+    "pack-refs": cmd_pack_refs,
     "pull": cmd_pull,
     "push": cmd_push,
     "receive-pack": cmd_receive_pack,

dulwich/porcelain.py

@@ -407,6 +407,18 @@ def symbolic_ref(repo, ref_name, force=False):
         repo_obj.refs.set_symbolic_ref(b"HEAD", ref_path)
 
 
+def pack_refs(repo, all=False):
+    with open_repo_closing(repo) as repo_obj:
+        refs = repo_obj.refs
+        packed_refs = {
+            ref: refs[ref]
+            for ref in refs
+            if (all or ref.startswith(LOCAL_TAG_PREFIX)) and ref != b"HEAD"
+        }
+
+        refs.add_packed_refs(packed_refs)
+
+
 def commit(
     repo=".",
     message=None,

dulwich/refs.py

@@ -155,6 +155,15 @@ class RefsContainer(object):
         """
         raise NotImplementedError(self.get_packed_refs)
 
+    def add_packed_refs(self, new_refs: Dict[Ref, Optional[ObjectID]]):
+        """Add the given refs as packed refs.
+
+        Args:
+          new_refs: A mapping of ref names to targets; if a target is None that
+            means remove the ref
+        """
+        raise NotImplementedError(self.add_packed_refs)
+
     def get_peeled(self, name):
         """Return the cached peeled value of a ref, if available.
 
@@ -706,6 +715,44 @@ class DiskRefsContainer(RefsContainer):
                         self._packed_refs[name] = sha
         return self._packed_refs
 
+    def add_packed_refs(self, new_refs: Dict[Ref, Optional[ObjectID]]):
+        """Add the given refs as packed refs.
+
+        Args:
+          new_refs: A mapping of ref names to targets; if a target is None that
+            means remove the ref
+        """
+        if not new_refs:
+            return
+
+        path = os.path.join(self.path, b"packed-refs")
+
+        with GitFile(path, "wb") as f:
+            # reread cached refs from disk, while holding the lock
+            packed_refs = self.get_packed_refs().copy()
+
+            for ref, target in new_refs.items():
+                # sanity check
+                if ref == HEADREF:
+                    raise ValueError("cannot pack HEAD")
+
+                # remove any loose refs pointing to this one -- please
+                # note that this bypasses remove_if_equals as we don't
+                # want to affect packed refs in here
+                try:
+                    os.remove(self.refpath(ref))
+                except (OSError, UnicodeError):
+                    pass
+
+                if target is not None:
+                    packed_refs[ref] = target
+                else:
+                    packed_refs.pop(ref, None)
+
+            write_packed_refs(f, packed_refs, self._peeled_refs)
+
+            self._packed_refs = packed_refs
+
     def get_peeled(self, name):
         """Return the cached peeled value of a ref, if available.
 

dulwich/tests/test_porcelain.py

@@ -3140,6 +3140,43 @@ class FindUniqueAbbrevTests(PorcelainTestCase):
             porcelain.find_unique_abbrev(self.repo.object_store, c1.id))
 
 
+class PackRefsTests(PorcelainTestCase):
+    def test_all(self):
+        c1, c2, c3 = build_commit_graph(
+            self.repo.object_store, [[1], [2, 1], [3, 1, 2]]
+        )
+        self.repo.refs[b"HEAD"] = c3.id
+        self.repo.refs[b"refs/heads/master"] = c2.id
+        self.repo.refs[b"refs/tags/foo"] = c1.id
+
+        porcelain.pack_refs(self.repo, all=True)
+
+        self.assertEqual(
+            self.repo.refs.get_packed_refs(),
+            {
+                b"refs/heads/master": c2.id,
+                b"refs/tags/foo": c1.id,
+            },
+        )
+
+    def test_not_all(self):
+        c1, c2, c3 = build_commit_graph(
+            self.repo.object_store, [[1], [2, 1], [3, 1, 2]]
+        )
+        self.repo.refs[b"HEAD"] = c3.id
+        self.repo.refs[b"refs/heads/master"] = c2.id
+        self.repo.refs[b"refs/tags/foo"] = c1.id
+
+        porcelain.pack_refs(self.repo)
+
+        self.assertEqual(
+            self.repo.refs.get_packed_refs(),
+            {
+                b"refs/tags/foo": c1.id,
+            },
+        )
+
+
 class ServerTests(PorcelainTestCase):
     @contextlib.contextmanager
     def _serving(self):

dulwich/tests/test_refs.py

@@ -447,6 +447,66 @@ class DiskRefsContainerTests(RefsContainerTests, TestCase):
             b"42d06bd4b77fed026b154d16493e5deab78f02ec",
         )
 
+        # this shouldn't overwrite the packed refs
+        self.assertEqual(
+            {b"refs/heads/packed": b"42d06bd4b77fed026b154d16493e5deab78f02ec"},
+            self._refs.get_packed_refs(),
+        )
+
+    def test_add_packed_refs(self):
+        # first, create a non-packed ref
+        self._refs[b"refs/heads/packed"] = b"3ec9c43c84ff242e3ef4a9fc5bc111fd780a76a8"
+
+        packed_ref_path = os.path.join(self._refs.path, b"refs", b"heads", b"packed")
+        self.assertTrue(os.path.exists(packed_ref_path))
+
+        # now overwrite that with a packed ref
+        packed_refs_file_path = os.path.join(self._refs.path, b"packed-refs")
+        self._refs.add_packed_refs(
+            {
+                b"refs/heads/packed": b"42d06bd4b77fed026b154d16493e5deab78f02ec",
+            }
+        )
+
+        # that should kill the file
+        self.assertFalse(os.path.exists(packed_ref_path))
+
+        # now delete the packed ref
+        self._refs.add_packed_refs(
+            {
+                b"refs/heads/packed": None,
+            }
+        )
+
+        # and it's gone!
+        self.assertFalse(os.path.exists(packed_ref_path))
+
+        self.assertRaises(
+            KeyError,
+            self._refs.__getitem__,
+            b"refs/heads/packed",
+        )
+
+        # just in case, make sure we can't pack HEAD
+        self.assertRaises(
+            ValueError,
+            self._refs.add_packed_refs,
+            {b"HEAD": "02ac81614bcdbd585a37b4b0edf8cb8a"},
+        )
+
+        # delete all packed refs
+        self._refs.add_packed_refs({ref: None for ref in self._refs.get_packed_refs()})
+
+        self.assertEqual({}, self._refs.get_packed_refs())
+
+        # remove the packed ref file, and check that adding nothing doesn't affect that
+        os.remove(packed_refs_file_path)
+
+        # adding nothing doesn't make it reappear
+        self._refs.add_packed_refs({})
+
+        self.assertFalse(os.path.exists(packed_refs_file_path))
+
     def test_setitem_symbolic(self):
         ones = b"1" * 40
         self._refs[b"HEAD"] = ones