Move find_shallow to object store.

dulwich/object_store.py

@@ -79,6 +79,47 @@ PACKDIR = "pack"
 PACK_MODE = 0o444 if sys.platform != "win32" else 0o644
 
 
+def find_shallow(store, heads, depth):
+    """Find shallow commits according to a given depth.
+
+    Args:
+      store: An ObjectStore for looking up objects.
+      heads: Iterable of head SHAs to start walking from.
+      depth: The depth of ancestors to include. A depth of one includes
+        only the heads themselves.
+    Returns: A tuple of (shallow, not_shallow), sets of SHAs that should be
+        considered shallow and unshallow according to the arguments. Note that
+        these sets may overlap if a commit is reachable along multiple paths.
+    """
+    parents = {}
+
+    def get_parents(sha):
+        result = parents.get(sha, None)
+        if not result:
+            result = store[sha].parents
+            parents[sha] = result
+        return result
+
+    todo = []  # stack of (sha, depth)
+    for head_sha in heads:
+        obj = store.peel_sha(head_sha)
+        if isinstance(obj, Commit):
+            todo.append((obj.id, 1))
+
+    not_shallow = set()
+    shallow = set()
+    while todo:
+        sha, cur_depth = todo.pop()
+        if cur_depth < depth:
+            not_shallow.add(sha)
+            new_depth = cur_depth + 1
+            todo.extend((p, new_depth) for p in get_parents(sha))
+        else:
+            shallow.add(sha)
+
+    return shallow, not_shallow
+
+
 def get_depth(
     store, head, get_parents=lambda commit: commit.parents, max_depth=None,
 ):

dulwich/repo.py

@@ -758,7 +758,7 @@ class BaseRepo(object):
         return self.object_store.peel_sha(self.refs[ref]).id
 
     def get_walker(self, include: Optional[List[bytes]] = None,
-                   *args, **kwargs):
+                   **kwargs):
         """Obtain a walker for this repository.
 
         Args:
@@ -791,7 +791,7 @@ class BaseRepo(object):
 
         kwargs["get_parents"] = lambda commit: self.get_parents(commit.id, commit)
 
-        return Walker(self.object_store, include, *args, **kwargs)
+        return Walker(self.object_store, include, **kwargs)
 
     def __getitem__(self, name: Union[ObjectID, Ref]):
         """Retrieve a Git object by SHA1 or ref.

dulwich/server.py

@@ -67,6 +67,8 @@ from dulwich.objects import (
     Commit,
     valid_hexsha,
 )
+from dulwich.object_store import find_shallow
+
 from dulwich.pack import (
     write_pack_objects,
 )
@@ -456,47 +458,6 @@ def _split_proto_line(line, allowed):
     raise GitProtocolError("Received invalid line from client: %r" % line)
 
 
-def _find_shallow(store, heads, depth):
-    """Find shallow commits according to a given depth.
-
-    Args:
-      store: An ObjectStore for looking up objects.
-      heads: Iterable of head SHAs to start walking from.
-      depth: The depth of ancestors to include. A depth of one includes
-        only the heads themselves.
-    Returns: A tuple of (shallow, not_shallow), sets of SHAs that should be
-        considered shallow and unshallow according to the arguments. Note that
-        these sets may overlap if a commit is reachable along multiple paths.
-    """
-    parents = {}
-
-    def get_parents(sha):
-        result = parents.get(sha, None)
-        if not result:
-            result = store[sha].parents
-            parents[sha] = result
-        return result
-
-    todo = []  # stack of (sha, depth)
-    for head_sha in heads:
-        obj = store.peel_sha(head_sha)
-        if isinstance(obj, Commit):
-            todo.append((obj.id, 1))
-
-    not_shallow = set()
-    shallow = set()
-    while todo:
-        sha, cur_depth = todo.pop()
-        if cur_depth < depth:
-            not_shallow.add(sha)
-            new_depth = cur_depth + 1
-            todo.extend((p, new_depth) for p in get_parents(sha))
-        else:
-            shallow.add(sha)
-
-    return shallow, not_shallow
-
-
 def _want_satisfied(store, haves, want, earliest):
     o = store[want]
     pending = collections.deque([o])
@@ -698,7 +659,7 @@ class _ProtocolGraphWalker(object):
             self.client_shallow.add(val)
         self.read_proto_line((None,))  # consume client's flush-pkt
 
-        shallow, not_shallow = _find_shallow(self.store, wants, depth)
+        shallow, not_shallow = find_shallow(self.store, wants, depth)
 
         # Update self.shallow instead of reassigning it since we passed a
         # reference to it before this method was called.

dulwich/tests/test_object_store.py

@@ -45,6 +45,7 @@ from dulwich.objects import (
     SubmoduleEncountered,
     S_IFGITLINK,
 )
+
 from dulwich.object_store import (
     DiskObjectStore,
     MemoryObjectStore,
@@ -53,6 +54,7 @@ from dulwich.object_store import (
     commit_tree_changes,
     read_packs_file,
     tree_lookup_path,
+    find_shallow,
 )
 from dulwich.pack import (
     REF_DELTA,
@@ -63,6 +65,7 @@ from dulwich.tests import (
     TestCase,
 )
 from dulwich.tests.utils import (
+    make_commit,
     make_object,
     make_tag,
     build_pack,
@@ -803,3 +806,90 @@ class TestReadPacksFile(TestCase):
                 )
             ),
         )
+
+
+class FindShallowTests(TestCase):
+    def setUp(self):
+        super(FindShallowTests, self).setUp()
+        self._store = MemoryObjectStore()
+
+    def make_commit(self, **attrs):
+        commit = make_commit(**attrs)
+        self._store.add_object(commit)
+        return commit
+
+    def make_linear_commits(self, n, message=b""):
+        commits = []
+        parents = []
+        for _ in range(n):
+            commits.append(self.make_commit(parents=parents, message=message))
+            parents = [commits[-1].id]
+        return commits
+
+    def assertSameElements(self, expected, actual):
+        self.assertEqual(set(expected), set(actual))
+
+    def test_linear(self):
+        c1, c2, c3 = self.make_linear_commits(3)
+
+        self.assertEqual(
+            (set([c3.id]), set([])), find_shallow(self._store, [c3.id], 1)
+        )
+        self.assertEqual(
+            (set([c2.id]), set([c3.id])),
+            find_shallow(self._store, [c3.id], 2),
+        )
+        self.assertEqual(
+            (set([c1.id]), set([c2.id, c3.id])),
+            find_shallow(self._store, [c3.id], 3),
+        )
+        self.assertEqual(
+            (set([]), set([c1.id, c2.id, c3.id])),
+            find_shallow(self._store, [c3.id], 4),
+        )
+
+    def test_multiple_independent(self):
+        a = self.make_linear_commits(2, message=b"a")
+        b = self.make_linear_commits(2, message=b"b")
+        c = self.make_linear_commits(2, message=b"c")
+        heads = [a[1].id, b[1].id, c[1].id]
+
+        self.assertEqual(
+            (set([a[0].id, b[0].id, c[0].id]), set(heads)),
+            find_shallow(self._store, heads, 2),
+        )
+
+    def test_multiple_overlapping(self):
+        # Create the following commit tree:
+        # 1--2
+        #  \
+        #   3--4
+        c1, c2 = self.make_linear_commits(2)
+        c3 = self.make_commit(parents=[c1.id])
+        c4 = self.make_commit(parents=[c3.id])
+
+        # 1 is shallow along the path from 4, but not along the path from 2.
+        self.assertEqual(
+            (set([c1.id]), set([c1.id, c2.id, c3.id, c4.id])),
+            find_shallow(self._store, [c2.id, c4.id], 3),
+        )
+
+    def test_merge(self):
+        c1 = self.make_commit()
+        c2 = self.make_commit()
+        c3 = self.make_commit(parents=[c1.id, c2.id])
+
+        self.assertEqual(
+            (set([c1.id, c2.id]), set([c3.id])),
+            find_shallow(self._store, [c3.id], 2),
+        )
+
+    def test_tag(self):
+        c1, c2 = self.make_linear_commits(2)
+        tag = make_tag(c2, name=b"tag")
+        self._store.add_object(tag)
+
+        self.assertEqual(
+            (set([c1.id]), set([c2.id])),
+            find_shallow(self._store, [tag.id], 2),
+        )

dulwich/tests/test_server.py

@@ -50,7 +50,6 @@ from dulwich.server import (
     PackHandler,
     _split_proto_line,
     serve_command,
-    _find_shallow,
     _ProtocolGraphWalker,
     ReceivePackHandler,
     SingleAckGraphWalkerImpl,
@@ -249,93 +248,6 @@ class UploadPackHandlerTestCase(TestCase):
         self.assertEqual([], self._handler.proto._received[1])
 
 
-class FindShallowTests(TestCase):
-    def setUp(self):
-        super(FindShallowTests, self).setUp()
-        self._store = MemoryObjectStore()
-
-    def make_commit(self, **attrs):
-        commit = make_commit(**attrs)
-        self._store.add_object(commit)
-        return commit
-
-    def make_linear_commits(self, n, message=b""):
-        commits = []
-        parents = []
-        for _ in range(n):
-            commits.append(self.make_commit(parents=parents, message=message))
-            parents = [commits[-1].id]
-        return commits
-
-    def assertSameElements(self, expected, actual):
-        self.assertEqual(set(expected), set(actual))
-
-    def test_linear(self):
-        c1, c2, c3 = self.make_linear_commits(3)
-
-        self.assertEqual(
-            (set([c3.id]), set([])), _find_shallow(self._store, [c3.id], 1)
-        )
-        self.assertEqual(
-            (set([c2.id]), set([c3.id])),
-            _find_shallow(self._store, [c3.id], 2),
-        )
-        self.assertEqual(
-            (set([c1.id]), set([c2.id, c3.id])),
-            _find_shallow(self._store, [c3.id], 3),
-        )
-        self.assertEqual(
-            (set([]), set([c1.id, c2.id, c3.id])),
-            _find_shallow(self._store, [c3.id], 4),
-        )
-
-    def test_multiple_independent(self):
-        a = self.make_linear_commits(2, message=b"a")
-        b = self.make_linear_commits(2, message=b"b")
-        c = self.make_linear_commits(2, message=b"c")
-        heads = [a[1].id, b[1].id, c[1].id]
-
-        self.assertEqual(
-            (set([a[0].id, b[0].id, c[0].id]), set(heads)),
-            _find_shallow(self._store, heads, 2),
-        )
-
-    def test_multiple_overlapping(self):
-        # Create the following commit tree:
-        # 1--2
-        #  \
-        #   3--4
-        c1, c2 = self.make_linear_commits(2)
-        c3 = self.make_commit(parents=[c1.id])
-        c4 = self.make_commit(parents=[c3.id])
-
-        # 1 is shallow along the path from 4, but not along the path from 2.
-        self.assertEqual(
-            (set([c1.id]), set([c1.id, c2.id, c3.id, c4.id])),
-            _find_shallow(self._store, [c2.id, c4.id], 3),
-        )
-
-    def test_merge(self):
-        c1 = self.make_commit()
-        c2 = self.make_commit()
-        c3 = self.make_commit(parents=[c1.id, c2.id])
-
-        self.assertEqual(
-            (set([c1.id, c2.id]), set([c3.id])),
-            _find_shallow(self._store, [c3.id], 2),
-        )
-
-    def test_tag(self):
-        c1, c2 = self.make_linear_commits(2)
-        tag = make_tag(c2, name=b"tag")
-        self._store.add_object(tag)
-
-        self.assertEqual(
-            (set([c1.id]), set([c2.id])),
-            _find_shallow(self._store, [tag.id], 2),
-        )
-
-
 class TestUploadPackHandler(UploadPackHandler):
     @classmethod
     def required_capabilities(self):

dulwich/walk.py

@@ -243,6 +243,7 @@ class Walker(object):
         self,
         store,
         include,
+        *,
         exclude=None,
         order=ORDER_DATE,
         reverse=False,