Traversable.traverse: Added as_edge option allowing to receive the source of the item as well to enable predicates to do more proper checking

Author: Byron

lib/git/objects/utils.py

@@ -92,7 +92,7 @@ def _get_intermediate_items(cls, item):
 
 	def traverse( self, predicate = lambda i,d: True,
 						   prune = lambda i,d: False, depth = -1, branch_first=True,
-						   visit_once = True, ignore_self=1 ):
+						   visit_once = True, ignore_self=1, as_edge = False ):
 		"""
 		``Returns``
 			iterator yieling of items found when traversing self
@@ -119,41 +119,49 @@ def traverse( self, predicate = lambda i,d: True,
 		
 		``ignore_self``
 			if True, self will be ignored and automatically pruned from
-			the result. Otherwise it will be the first item to be returned"""
+			the result. Otherwise it will be the first item to be returned.
+			If as_edge is True, the source of the first edge is None
+			
+		``as_edge``
+			if True, return a pair of items, first being the source, second the 
+			destinatination, i.e. tuple(src, dest) with the edge spanning from 
+			source to destination"""
 		visited = set()
 		stack = Deque()
-		stack.append( ( 0 ,self ) )		# self is always depth level 0
+		stack.append( ( 0 ,self, None ) )		# self is always depth level 0
 
-		def addToStack( stack, lst, branch_first, dpth ):
+		def addToStack( stack, item, branch_first, depth ):
+			lst = self._get_intermediate_items( item )
 			if not lst:
 				return
 			if branch_first:
-				stack.extendleft( ( dpth , item ) for item in lst )
+				stack.extendleft( ( depth , i, item ) for i in lst )
 			else:
-				reviter = ( ( dpth , lst[i] ) for i in range( len( lst )-1,-1,-1) )
+				reviter = ( ( depth , lst[i], item ) for i in range( len( lst )-1,-1,-1) )
 				stack.extend( reviter )
 		# END addToStack local method
 
 		while stack:
-			d, item = stack.pop()			# depth of item, item
+			d, item, src = stack.pop()			# depth of item, item, item_source
 
 			if visit_once and item in visited:
 				continue
 
 			if visit_once:
 				visited.add(item)
 
-			if prune( item, d ):
+			rval = ( as_edge and (src, item) ) or item
+			if prune( rval, d ):
 				continue
 
 			skipStartItem = ignore_self and ( item == self )
-			if not skipStartItem and predicate( item, d ):
-				yield item
+			if not skipStartItem and predicate( rval, d ):
+				yield rval
 
 			# only continue to next level if this is appropriate !
 			nd = d + 1
 			if depth > -1 and nd > depth:
 				continue
 
-			addToStack( stack, self._get_intermediate_items( item ), branch_first, nd )
+			addToStack( stack, item, branch_first, nd )
 		# END for each item on work stack

test/git/test_commit.py

@@ -69,8 +69,9 @@ def test_traversal(self):
 
 		# at some point, both iterations should stop
 		assert list(bfirst)[-1] == first
-		stoptraverse = self.rorepo.commit("254d04aa3180eb8b8daf7b7ff25f010cd69b4e7d").traverse()
-		list(stoptraverse)
+		stoptraverse = self.rorepo.commit("254d04aa3180eb8b8daf7b7ff25f010cd69b4e7d").traverse(as_edge=True)
+		l = list(stoptraverse)
+		assert len(l[0]) == 2
 
 		# ignore self
 		assert start.traverse(ignore_self=False).next() == start