walk: Walk a fixed number of extra commits before returning.

Change-Id: Ie93f5733e323e2acfbfc88a255c7e6ff8ab33d8c

dulwich/tests/test_walk.py

@@ -95,6 +95,7 @@ class WalkerTest(TestCase):
 
     def assertWalkYields(self, expected, *args, **kwargs):
         walker = Walker(self.store, *args, **kwargs)
+        expected = list(expected)
         for i, entry in enumerate(expected):
             if isinstance(entry, Commit):
                 expected[i] = TestWalkEntry(entry, None)
@@ -111,24 +112,17 @@ class WalkerTest(TestCase):
         self.assertWalkYields([c3, c2], [c3.id, c1.id], exclude=[c1.id])
         self.assertWalkYields([c3], [c3.id, c1.id], exclude=[c2.id])
 
-    def assertNextFails(self, walk_iter, missing_sha):
-        try:
-            next(walk_iter)
-            self.fail('Failed to error on missing sha %s' % missing_sha)
-        except MissingCommitError, e:
-            self.assertEqual(missing_sha, e.sha)
-
     def test_missing(self):
-        c1, c2, c3 = self.make_linear_commits(3)
-        self.assertWalkYields([c3, c2, c1], [c3.id])
-
-        del self.store[c1.id]
-        self.assertWalkYields([c3], [c3.id], max_entries=1)
-        walk_iter = iter(Walker(self.store, [c3.id]))
-        self.assertEqual(TestWalkEntry(c3, None), next(walk_iter))
-        self.assertNextFails(walk_iter, c1.id)
-        self.assertNextFails(iter(Walker(self.store, [c2.id])), c1.id)
-        self.assertRaises(MissingCommitError, Walker, self.store, c1.id)
+        cs = list(reversed(self.make_linear_commits(20)))
+        self.assertWalkYields(cs, [cs[0].id])
+
+        # Exactly how close we can get to a missing commit depends on our
+        # implementation (in particular the choice of _MAX_EXTRA_COMMITS), but
+        # we should at least be able to walk some history in a broken repo.
+        del self.store[cs[-1].id]
+        for i in xrange(1, 11):
+            self.assertWalkYields(cs[:i], [cs[0].id], max_entries=i)
+        self.assertRaises(MissingCommitError, Walker, self.store, cs[0].id)
 
     def test_branch(self):
         c1, x2, x3, y4 = self.make_commits([[1], [2, 1], [3, 2], [4, 1]])

dulwich/walk.py

@@ -98,6 +98,7 @@ class _CommitTimeQueue(object):
         self._done = set()
         self._min_time = walker.since
         self._extra_commits_left = _MAX_EXTRA_COMMITS
+        self._is_finished = False
 
         for commit_id in itertools.chain(walker.include, walker.excluded):
             self._push(commit_id)
@@ -112,6 +113,8 @@ class _CommitTimeQueue(object):
             self._pq_set.add(commit_id)
 
     def next(self):
+        if self._is_finished:
+            return None
         while self._pq:
             _, commit = heapq.heappop(self._pq)
             sha = commit.id
@@ -142,6 +145,7 @@ class _CommitTimeQueue(object):
 
             if not is_excluded:
                 return WalkEntry(self._walker, commit)
+        self._is_finished = True
         return None
 
 
@@ -198,6 +202,7 @@ class Walker(object):
 
         self._num_entries = 0
         self._queue = queue_cls(self)
+        self._out_queue = collections.deque()
 
     def _path_matches(self, changed_path):
         if changed_path is None:
@@ -254,15 +259,18 @@ class Walker(object):
 
     def _next(self):
         max_entries = self.max_entries
-        while True:
-            if max_entries is not None and self._num_entries >= max_entries:
-                return None
+        while max_entries is None or self._num_entries < max_entries:
             entry = self._queue.next()
-            if entry is None:
-                return None
-            if self._should_return(entry):
-                self._num_entries += 1
-                return entry
+            if entry is not None:
+                self._out_queue.append(entry)
+            if entry is None or len(self._out_queue) > _MAX_EXTRA_COMMITS:
+                if not self._out_queue:
+                    return None
+                entry = self._out_queue.popleft()
+                if self._should_return(entry):
+                    self._num_entries += 1
+                    return entry
+        return None
 
     def _reorder(self, results):
         """Possibly reorder a results iterator.