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

Author: Byron

lib/git/index.py

@@ -284,7 +284,7 @@ def from_tree(cls, repo, *treeish, **kwargs):
 			
 		``*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 @@ def from_tree(cls, repo, *treeish, **kwargs):
 			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))

lib/git/objects/tree.py

@@ -152,22 +152,21 @@ def __repr__(self):
 		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 @@ def traverse(self, max_depth=-1, predicate = lambda i: True):
 		``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):

test/git/test_tree.py

@@ -35,6 +35,11 @@ def test_traverse(self):
 		trees = list(root.traverse(predicate = trees_only))
 		assert len(trees) == len(list( i for i in root.traverse() if trees_only(i) ))
 
+		# test prune
+		lib_folder = lambda t: t.path == "lib"
+		pruned_trees = list(root.traverse(predicate = trees_only,prune = lib_folder))
+		assert len(pruned_trees) < len(trees)
+		
 		# trees and blobs
 		assert len(set(trees)|set(root.trees)) == len(trees)
 		assert len(set(b for b in root if isinstance(b, Blob)) | set(root.blobs)) == len( root.blobs )