dulwich-users team mailing list archive

[dborowitz@xxxxxxxxxx]

From: Dave Borowitz <dborowitz@xxxxxxxxxx>

This allows us to return data from walks in addition to just the commit.

Change-Id: Ied201689ed7621118235a85206debaefb5228e49
---
 dulwich/tests/test_walk.py |   23 +++++++++++++++--------
 dulwich/walk.py            |   23 ++++++++++++++++++-----
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/dulwich/tests/test_walk.py b/dulwich/tests/test_walk.py
index 476ab2a..ac9d964 100644
--- a/dulwich/tests/test_walk.py
+++ b/dulwich/tests/test_walk.py
@@ -21,7 +21,11 @@
 from dulwich.object_store import (
     MemoryObjectStore,
     )
+from dulwich.objects import (
+    Commit,
+    )
 from dulwich.walk import (
+    WalkEntry,
     Walker,
     )
 from dulwich.tests import TestCase
@@ -49,6 +53,9 @@ class WalkerTest(TestCase):
 
     def assertWalkYields(self, expected, *args, **kwargs):
         walker = Walker(self.store, *args, **kwargs)
+        for i, entry in enumerate(expected):
+            if isinstance(entry, Commit):
+                expected[i] = WalkEntry(entry)
         actual = list(walker)
         self.assertEqual(expected, actual)
 
@@ -76,15 +83,15 @@ class WalkerTest(TestCase):
         c1, c2, c3 = self.make_linear_commits(3)
         self.assertWalkYields([c1, c2, c3], [c3.id], reverse=True)
 
-    def test_max_commits(self):
+    def test_max_entries(self):
         c1, c2, c3 = self.make_linear_commits(3)
-        self.assertWalkYields([c3, c2, c1], [c3.id], max_commits=3)
-        self.assertWalkYields([c3, c2], [c3.id], max_commits=2)
-        self.assertWalkYields([c3], [c3.id], max_commits=1)
+        self.assertWalkYields([c3, c2, c1], [c3.id], max_entries=3)
+        self.assertWalkYields([c3, c2], [c3.id], max_entries=2)
+        self.assertWalkYields([c3], [c3.id], max_entries=1)
 
-    def test_reverse_after_max_commits(self):
+    def test_reverse_after_max_entries(self):
         c1, c2, c3 = self.make_linear_commits(3)
-        self.assertWalkYields([c1, c2, c3], [c3.id], max_commits=3,
+        self.assertWalkYields([c1, c2, c3], [c3.id], max_entries=3,
                               reverse=True)
-        self.assertWalkYields([c2, c3], [c3.id], max_commits=2, reverse=True)
-        self.assertWalkYields([c3], [c3.id], max_commits=1, reverse=True)
+        self.assertWalkYields([c2, c3], [c3.id], max_entries=2, reverse=True)
+        self.assertWalkYields([c3], [c3.id], max_entries=1, reverse=True)
diff --git a/dulwich/walk.py b/dulwich/walk.py
index ce123d1..3a84864 100644
--- a/dulwich/walk.py
+++ b/dulwich/walk.py
@@ -24,6 +24,16 @@ import itertools
 ORDER_DATE = 'date'
 
 
+class WalkEntry(object):
+    """Object encapsulating a single result from a walk."""
+
+    def __init__(self, commit):
+        self.commit = commit
+
+    def __eq__(self, other):
+        return isinstance(other, WalkEntry) and self.commit == other.commit
+
+
 class Walker(object):
     """Object for performing a walk of commits in a store.
 
@@ -32,7 +42,7 @@ class Walker(object):
     """
 
     def __init__(self, store, include, exclude=None, order=ORDER_DATE,
-                 reverse=False, max_commits=None):
+                 reverse=False, max_entries=None):
         """Constructor.
 
         :param store: ObjectStore instance for looking up objects.
@@ -44,7 +54,7 @@ class Walker(object):
             other than ORDER_DATE may result in O(n) memory usage.
         :param reverse: If True, reverse the order of output, requiring O(n)
             memory.
-        :param max_commits: The maximum number of commits to yield, or None for
+        :param max_entries: The maximum number of entries to yield, or None for
             no limit.
         """
         self._store = store
@@ -53,7 +63,7 @@ class Walker(object):
             raise ValueError('Unknown walk order %s' % order)
         self._order = order
         self._reverse = reverse
-        self._max_commits = max_commits
+        self._max_entries = max_entries
 
         exclude = exclude or []
         self._excluded = set(exclude)
@@ -91,10 +101,13 @@ class Walker(object):
         return None
 
     def _next(self):
-        limit = self._max_commits
+        limit = self._max_entries
         if limit is not None and len(self._done) >= limit:
             return None
-        return self._pop()
+        commit = self._pop()
+        if commit is None:
+            return None
+        return WalkEntry(commit)
 
     def __iter__(self):
         results = iter(self._next, None)
-- 
1.7.3.1