index.iter_blobs method added including tests ( which have been improved generally for more coverage )

2 files changed, 71 insertions, 12 deletions

diff --git a/lib/git/index.py b/lib/git/index.py
index 7481f4ce..e5b2d009 100644
--- a/lib/git/index.py
+++ b/lib/git/index.py
@@ -13,6 +13,8 @@ import mmap
 import objects
 import tempfile
 import os
+import stat
+from git.objects import Blob
 
 class IndexEntry(tuple):
 	"""
@@ -114,19 +116,24 @@ class Index(object):
 	Implements an Index that can be manipulated using a native implementation in 
 	order to save git command function calls wherever possible.
 	
-	It provides custom merging facilities and to create custom commits.
+	It provides custom merging facilities allowing to merge without actually changing
+	your index or your working tree. This way you can perform own test-merges based
+	on the index only without having to deal with the working copy. This is useful 
+	in case of partial working trees.
 	
 	``Entries``
 	The index contains an entries dict whose keys are tuples of type IndexEntry
 	to facilitate access.
 	"""
-	__slots__ = ( "version", "entries", "_extension_data" )
+	__slots__ = ( "repo", "version", "entries", "_extension_data" )
 	_VERSION = 2			# latest version we support
+	S_IFGITLINK	= 0160000
 	
-	def __init__(self, stream = None):
+	def __init__(self, repo, stream = None):
 		"""
 		Initialize this Index instance, optionally from the given ``stream``
 		"""
+		self.repo = repo
 		self.entries = dict()
 		self.version = self._VERSION
 		self._extension_data = ''
@@ -178,11 +185,14 @@ class Index(object):
 		self._extension_data = stream.read(~0)
 	
 	@classmethod
-	def from_file(cls, file_path):
+	def from_file(cls, repo, file_path):
 		"""
 		Returns
 			Index instance as recreated from the given stream.
-			
+		
+		``repo``
+			Repository the index is related to
+		
 		``file_pa ``
 			File path pointing to git index file
 			
@@ -200,7 +210,7 @@ class Index(object):
 		# END memory mapping
 		
 		try:
-			return cls(stream)
+			return cls(repo, stream)
 		finally:
 			fp.close()
 		
@@ -331,7 +341,7 @@ class Index(object):
 		try:
 			os.rename(cur_index, moved_index)
 			repo.git.read_tree(*arg_list, **kwargs)
-			index = cls.from_file(tmp_index)
+			index = cls.from_file(repo, tmp_index)
 		finally:
 			# put back the original index first !
 			if os.path.exists(moved_index):
@@ -341,7 +351,41 @@ class Index(object):
 		# END index merge handling
 		
 		return index
-		
+	
+	@classmethod
+	def _index_mode_to_tree_index_mode(cls, index_mode):
+		"""Cleanup a index_mode value.
+		This will return a index_mode that can be stored in a tree object.
+		``index_mode``
+			Index_mode to clean up.
+		"""
+		if stat.S_ISLNK(index_mode):
+			return stat.S_IFLNK
+		elif stat.S_ISDIR(index_mode):
+			return stat.S_IFDIR
+		elif stat.S_IFMT(index_mode) == cls.S_IFGITLINK:
+			return cls.S_IFGITLINK
+		ret = stat.S_IFREG | 0644
+		ret |= (index_mode & 0111)
+		return ret
+	
+	def iter_blobs(self, predicate = lambda t: True):
+		"""
+		Returns
+			Iterator yielding tuples of Blob objects and stages, tuple(stage, Blob)
+			
+		``predicate``
+			Function(t) returning True if tuple(stage, Blob) should be yielded by the 
+			iterator
+		"""
+		for entry in self.entries.itervalues():
+			mode = self._index_mode_to_tree_index_mode(entry.mode)
+			blob = Blob(self.repo, entry.sha, mode, entry.path)
+			output = (entry.stage, blob)
+			if predicate(output):
+				yield output
+		# END for each entry 
+	
 	def write_tree(self, stream):
 		"""
 		Writes the 
diff --git a/test/git/test_index.py b/test/git/test_index.py
index d256e7c0..524f0778 100644
--- a/test/git/test_index.py
+++ b/test/git/test_index.py
@@ -8,6 +8,7 @@ from test.testlib import *
 from git import *
 import inspect
 import os
+import tempfile
 
 class TestTree(TestCase):
 	
@@ -17,7 +18,7 @@ class TestTree(TestCase):
 		
 	def test_base(self):
 		# read from file
-		index = Index.from_file(fixture_path("index"))
+		index = Index.from_file(self.repo, fixture_path("index"))
 		assert index.entries
 		assert index.version > 0
 		
@@ -30,7 +31,7 @@ class TestTree(TestCase):
 		# END for each method
 		
 		# test stage
-		index_merge = Index.from_file(fixture_path("index_merge"))
+		index_merge = Index.from_file(self.repo, fixture_path("index_merge"))
 		assert len(index_merge.entries) == 106
 		assert len(list(e for e in index_merge.entries.itervalues() if e.stage != 0 ))
 		
@@ -40,6 +41,11 @@ class TestTree(TestCase):
 		
 		index_output.seek(0)
 		assert index_output.read() == fixture("index_merge")
+		
+		tmpfile = tempfile.mktemp()
+		Index.to_file(index_merge, tmpfile)
+		assert os.path.isfile(tmpfile)
+		os.remove(tmpfile)
 	
 	def _cmp_tree_index(self, tree, index):
 		# fail unless both objects contain the same paths and blobs
@@ -53,7 +59,7 @@ class TestTree(TestCase):
 		# END for each blob in tree
 		assert num_blobs == len(index.entries)
 	
-	def test_merge(self):
+	def test_from_tree(self):
 		common_ancestor_sha = "5117c9c8a4d3af19a9958677e45cda9269de1541"
 		cur_sha = "4b43ca7ff72d5f535134241e7c797ddc9c7a3573"
 		other_sha = "39f85c4358b7346fee22169da9cad93901ea9eb9"
@@ -70,7 +76,16 @@ class TestTree(TestCase):
 		
 		# merge three trees - here we have a merge conflict
 		tree_way_index = Index.from_tree(self.repo, common_ancestor_sha, cur_sha, other_sha)
-		assert len(list(e for e in tree_way_index.entries.values() if e.stage != 0)) 
+		assert len(list(e for e in tree_way_index.entries.values() if e.stage != 0))
+		
+		
+		# ITERATE BLOBS
+		merge_required = lambda t: t[0] != 0
+		merge_blobs = list(tree_way_index.iter_blobs(merge_required))
+		assert merge_blobs
+		assert merge_blobs[0][0] in (1,2,3)
+		assert isinstance(merge_blobs[0][1], Blob)
+		
 		
 	def test_custom_commit(self):
 		self.fail("Custom commit:write tree, make commit with custom parents")