ItemTraversal: Predicate and prune functions now provide depth information, allowing the callee to know more about its environment

4 files changed, 13 insertions, 13 deletions

diff --git a/lib/git/objects/utils.py b/lib/git/objects/utils.py
index 18968c0f..a3c6edc7 100644
--- a/lib/git/objects/utils.py
+++ b/lib/git/objects/utils.py
@@ -90,18 +90,18 @@ class Traversable(object):
 		raise NotImplementedError("To be implemented in subclass")
 			
 	
-	def traverse( self, predicate = lambda i: True,
-						   prune = lambda i: False, depth = -1, branch_first=True,
+	def traverse( self, predicate = lambda i,d: True,
+						   prune = lambda i,d: False, depth = -1, branch_first=True,
 						   visit_once = True, ignore_self=1 ):
 		"""
 		``Returns``
 			iterator yieling of items found when traversing self
 			
 		``predicate``
-			f(i) returns False if item i should not be included in the result
+			f(i,d) returns False if item i at depth d should not be included in the result
 			
 		``prune``
-			f(i) return True if the search should stop at item i.
+			f(i,d) return True if the search should stop at item i at depth d.
 			Item i will not be returned.
 			
 		``depth``
@@ -132,11 +132,11 @@ class Traversable(object):
 		while stack:
 			d, item = stack.pop()			# depth of item, item
 			
-			if prune( item ):
+			if prune( item, d ):
 				continue
 	
 			skipStartItem = ignore_self and ( item == self )
-			if not skipStartItem and predicate( item ):
+			if not skipStartItem and predicate( item, d ):
 				yield item
 	
 			# only continue to next level if this is appropriate !
diff --git a/test/git/test_commit.py b/test/git/test_commit.py
index e06e35f4..0c05ba50 100644
--- a/test/git/test_commit.py
+++ b/test/git/test_commit.py
@@ -69,10 +69,10 @@ class TestCommit(TestBase):
 		assert len(list(start.traverse(ignore_self=False, depth=0))) == 1
 		
 		# prune
-		assert start.traverse(branch_first=1, prune=lambda i: i==p0).next() == p1
+		assert start.traverse(branch_first=1, prune=lambda i,d: i==p0).next() == p1
 		
 		# predicate
-		assert start.traverse(branch_first=1, predicate=lambda i: i==p1).next() == p1
+		assert start.traverse(branch_first=1, predicate=lambda i,d: i==p1).next() == p1
 		
 	@patch_object(Git, '_call_process')
 	def test_rev_list_bisect_all(self, git):
diff --git a/test/git/test_index.py b/test/git/test_index.py
index 2f4da161..e96aff4c 100644
--- a/test/git/test_index.py
+++ b/test/git/test_index.py
@@ -84,7 +84,7 @@ class TestTree(TestBase):
 		
 		num_blobs = 0
 		blist = list()
-		for blob in tree.traverse(predicate = lambda e: e.type == "blob", branch_first=False):
+		for blob in tree.traverse(predicate = lambda e,d: e.type == "blob", branch_first=False):
 			assert (blob.path,0) in index.entries
 			blist.append(blob)
 		# END for each blob in tree
@@ -134,7 +134,7 @@ class TestTree(TestBase):
 		tree = three_way_index.write_tree()
 		assert isinstance(tree, Tree)
 		num_blobs = 0
-		for blob in tree.traverse(predicate=lambda item: item.type == "blob"):
+		for blob in tree.traverse(predicate=lambda item,d: item.type == "blob"):
 			assert (blob.path,0) in three_way_index.entries
 			num_blobs += 1
 		# END for each blob
diff --git a/test/git/test_tree.py b/test/git/test_tree.py
index 2bfd9922..0bb9647f 100644
--- a/test/git/test_tree.py
+++ b/test/git/test_tree.py
@@ -32,12 +32,12 @@ class TestTree(TestCase):
 		assert len(list(root)) == len(list(root.traverse(depth=1)))
 		
 		# only choose trees
-		trees_only = lambda i: i.type == "tree"
+		trees_only = lambda i,d: i.type == "tree"
 		trees = list(root.traverse(predicate = trees_only))
-		assert len(trees) == len(list( i for i in root.traverse() if trees_only(i) ))
+		assert len(trees) == len(list( i for i in root.traverse() if trees_only(i,0) ))
 		
 		# test prune
-		lib_folder = lambda t: t.path == "lib"
+		lib_folder = lambda t,d: t.path == "lib"
 		pruned_trees = list(root.traverse(predicate = trees_only,prune = lib_folder))
 		assert len(pruned_trees) < len(trees)