Merge pull request #1049 from jelmer/aiohttp-refactor

Factor out some formatting functions ahead of aiohttp support

dulwich/protocol.py

@@ -585,3 +585,17 @@ class PktLineParser(object):
     def get_tail(self):
         """Read back any unused data."""
         return self._readahead.getvalue()
+
+
+def format_capability_line(capabilities):
+    return b"".join([b" " + c for c in capabilities])
+
+
+def format_ref_line(ref, sha, capabilities=None):
+    if capabilities is None:
+        return sha + b" " + ref + b"\n"
+    else:
+        return (
+            sha + b" " + ref + b"\0"
+            + format_capability_line(capabilities)
+            + b"\n")

dulwich/server.py

@@ -108,6 +108,7 @@ from dulwich.protocol import (
     extract_capabilities,
     extract_want_line_capabilities,
     symref_capabilities,
+    format_ref_line,
 )
 from dulwich.refs import (
     ANNOTATED_TAG_SUFFIX,
@@ -233,11 +234,6 @@ class PackHandler(Handler):
         # Flags needed for the no-done capability
         self._done_received = False
 
-    @classmethod
-    def capability_line(cls, capabilities):
-        logger.info("Sending capabilities: %s", capabilities)
-        return b"".join([b" " + c for c in capabilities])
-
     @classmethod
     def capabilities(cls) -> Iterable[bytes]:
         raise NotImplementedError(cls.capabilities)
@@ -602,17 +598,19 @@ class _ProtocolGraphWalker(object):
                     # TODO(jelmer): Integrate with Repo.fetch_objects refs
                     # logic.
                     continue
-                line = sha + b" " + ref
-                if not i:
-                    line += b"\x00" + self.handler.capability_line(
+                if i == 0:
+                    logger.info(
+                        "Sending capabilities: %s", self.handler.capabilities())
+                    line = format_ref_line(
+                        ref, sha,
                         self.handler.capabilities()
-                        + symref_capabilities(symrefs.items())
-                    )
-                self.proto.write_pkt_line(line + b"\n")
+                        + symref_capabilities(symrefs.items()))
+                else:
+                    line = format_ref_line(ref, sha)
+                self.proto.write_pkt_line(line)
                 if peeled_sha != sha:
                     self.proto.write_pkt_line(
-                        peeled_sha + b" " + ref + ANNOTATED_TAG_SUFFIX + b"\n"
-                    )
+                        format_ref_line(ref + ANNOTATED_TAG_SUFFIX, peeled_sha))
 
             # i'm done..
             self.proto.write_pkt_line(None)
@@ -1047,19 +1045,15 @@ class ReceivePackHandler(PackHandler):
 
             if not refs:
                 refs = [(CAPABILITIES_REF, ZERO_SHA)]
+            logger.info(
+                "Sending capabilities: %s", self.capabilities())
             self.proto.write_pkt_line(
-                refs[0][1]
-                + b" "
-                + refs[0][0]
-                + b"\0"
-                + self.capability_line(
-                    self.capabilities() + symref_capabilities(symrefs)
-                )
-                + b"\n"
-            )
+                format_ref_line(
+                    refs[0][0], refs[0][1],
+                    self.capabilities() + symref_capabilities(symrefs)))
             for i in range(1, len(refs)):
                 ref = refs[i]
-                self.proto.write_pkt_line(ref[1] + b" " + ref[0] + b"\n")
+                self.proto.write_pkt_line(format_ref_line(ref[0], ref[1]))
 
             self.proto.write_pkt_line(None)
             if self.advertise_refs:

dulwich/tests/test_server.py

@@ -64,6 +64,7 @@ from dulwich.tests.utils import (
 )
 from dulwich.protocol import (
     ZERO_SHA,
+    format_capability_line,
 )
 
 ONE = b"1" * 40
@@ -131,7 +132,7 @@ class HandlerTestCase(TestCase):
     def test_capability_line(self):
         self.assertEqual(
             b" cap1 cap2 cap3",
-            self._handler.capability_line([b"cap1", b"cap2", b"cap3"]),
+            format_capability_line([b"cap1", b"cap2", b"cap3"]),
         )
 
     def test_set_client_capabilities(self):