tree.traverse: Added prune functionality - previously the predciate did both, pruning and preventing to return items

2 files changed, 17 insertions, 9 deletions

diff --git a/lib/git/index.py b/lib/git/index.py
index 1042d7b8..7481f4ce 100644
--- a/lib/git/index.py
+++ b/lib/git/index.py
@@ -284,7 +284,7 @@ class Index(object):
 			
 		``*treeish``
 			One, two or three Tree Objects or Commits. The result changes according to the 
-			amoutn of trees.
+			amount of trees.
 			If 1 Tree is given, it will just be read into a new index
 			If 2 Trees are given, they will be merged into a new index using a 
 			 two way merge algorithm. Tree 1 is the 'current' tree, tree 2 is the 'other'
@@ -300,6 +300,10 @@ class Index(object):
 			In the three-way merge case, --aggressive will be specified to automatically
 			resolve more cases in a commonly correct manner. Specify trivial=True as kwarg
 			to override that.
+			
+			As the underlying git-read-tree command takes into account the current index, 
+			it will be temporarily moved out of the way to assure there are no unsuspected
+			interferences.
 		"""
 		if len(treeish) == 0 or len(treeish) > 3:
 			raise ValueError("Please specify between 1 and 3 treeish, got %i" % len(treeish))
diff --git a/lib/git/objects/tree.py b/lib/git/objects/tree.py
index c35c075e..92aae881 100644
--- a/lib/git/objects/tree.py
+++ b/lib/git/objects/tree.py
@@ -152,22 +152,21 @@ class Tree(base.IndexObject, diff.Diffable):
 		return '<git.Tree "%s">' % self.id
 		
 	@classmethod
-	def _iter_recursive(cls, repo, tree, cur_depth, max_depth, predicate ):
+	def _iter_recursive(cls, repo, tree, cur_depth, max_depth, predicate, prune ):
 		
 		for obj in tree:
 			# adjust path to be complete
 			obj.path = os.path.join(tree.path, obj.path)
-			if not predicate(obj):
-				continue
-			yield obj
-			if obj.type == "tree" and ( max_depth < 0 or cur_depth+1 <= max_depth ):
-				for recursive_obj in cls._iter_recursive( repo, obj, cur_depth+1, max_depth, predicate ):
+			if predicate(obj):
+				yield obj
+			if obj.type == "tree" and ( max_depth < 0 or cur_depth+1 <= max_depth ) and not prune(obj):
+				for recursive_obj in cls._iter_recursive( repo, obj, cur_depth+1, max_depth, predicate, prune ):
 					yield recursive_obj
 				# END for each recursive object
 			# END if we may enter recursion
 		# END for each object
 		
-	def traverse(self, max_depth=-1, predicate = lambda i: True):
+	def traverse(self, max_depth=-1, predicate = lambda i: True, prune = lambda t: False):
 		"""
 		Returns
 		
@@ -183,8 +182,13 @@ class Tree(base.IndexObject, diff.Diffable):
 		``predicate``
 		
 			If predicate(item) returns True, item will be returned by iterator
+			
+		``prune``
+			
+			If prune(tree) returns True, the traversal will not continue into the 
+			given tree object.
 		"""
-		return self._iter_recursive( self.repo, self, 0, max_depth, predicate )
+		return self._iter_recursive( self.repo, self, 0, max_depth, predicate, prune )
 		
 	@property
 	def trees(self):