tree: added TreeModifier, allowing to adjust existing trees safely and or fast, while staying compatible with serialization which requires it to be sorted

5 files changed, 152 insertions, 16 deletions

diff --git a/lib/git/objects/__init__.py b/lib/git/objects/__init__.py
index 717fa808..ef6b2ea4 100644
--- a/lib/git/objects/__init__.py
+++ b/lib/git/objects/__init__.py
@@ -7,6 +7,7 @@ from tag import *
 from blob import *
 from tree import *
 from commit import *
+from submodule import *
 
 __all__ = [ name for name, obj in locals().items()
             if not (name.startswith('_') or inspect.ismodule(obj)) ]
\ No newline at end of file
diff --git a/lib/git/objects/base.py b/lib/git/objects/base.py
index 5a3a15a7..36bbccb2 100644
--- a/lib/git/objects/base.py
+++ b/lib/git/objects/base.py
@@ -197,7 +197,7 @@ class IndexObject(Object):
 			for example.
 		"""
 		mode = 0
-		for iteration,char in enumerate(reversed(modestr[-6:])):
+		for iteration, char in enumerate(reversed(modestr[-6:])):
 			mode += int(char) << iteration*3
 		# END for each char
 		return mode
diff --git a/lib/git/objects/commit.py b/lib/git/objects/commit.py
index 0d8f2c64..3b20e314 100644
--- a/lib/git/objects/commit.py
+++ b/lib/git/objects/commit.py
@@ -402,7 +402,7 @@ class Commit(base.Object, Iterable, diff.Diffable, utils.Traversable, utils.Seri
 		""":param from_rev_list: if true, the stream format is coming from the rev-list command
 		Otherwise it is assumed to be a plain data stream from our object"""
 		readline = stream.readline
-		self.tree = Tree(self.repo, readline().split()[1], 0, '')
+		self.tree = Tree(self.repo, readline().split()[1], Tree.tree_id<<12, '')
 
 		self.parents = list()
 		next_line = None
diff --git a/lib/git/objects/tree.py b/lib/git/objects/tree.py
index 67aea1cb..0da62060 100644
--- a/lib/git/objects/tree.py
+++ b/lib/git/objects/tree.py
@@ -21,6 +21,84 @@ def sha_to_hex(sha):
 	return hexsha
 	
 
+class TreeModifier(object):
+	"""A utility class providing methods to alter the underlying cache in a list-like
+	fashion.
+	Once all adjustments are complete, the _cache, which really is a refernce to 
+	the cache of a tree, will be sorted. Assuring it will be in a serializable state"""
+	__slots__ = '_cache'
+	
+	def __init__(self, cache):
+		self._cache = cache
+	
+	def _index_by_name(self, name):
+		""":return: index of an item with name, or -1 if not found"""
+		for i, t in enumerate(self._cache):
+			if t[2] == name:
+				return i
+			# END found item
+		# END for each item in cache
+		return -1
+	
+	#{ Interface 
+	def set_done(self):
+		"""Call this method once you are done modifying the tree information.
+		It may be called several times, but be aware that each call will cause 
+		a sort operation
+		:return self:"""
+		self._cache.sort(key=lambda t: t[2])	# sort by name
+		return self
+	#} END interface
+	
+	#{ Mutators
+	def add(self, hexsha, mode, name, force=False):
+		"""Add the given item to the tree. If an item with the given name already
+		exists, nothing will be done, but a ValueError will be raised if the 
+		sha and mode of the existing item do not match the one you add, unless 
+		force is True
+		:param hexsha: The 40 byte sha of the item to add
+		:param mode: int representing the stat compatible mode of the item
+		:param force: If True, an item with your name and information will overwrite
+			any existing item with the same name, no matter which information it has
+		:return: self"""
+		if '/' in name:
+			raise ValueError("Name must not contain '/' characters")
+		if len(hexsha) != 40:
+			raise ValueError("Hexsha required, got %r" % hexsha)
+		if (mode >> 12) not in Tree._map_id_to_type:
+			raise ValueError("Invalid object type according to mode %o" % mode)
+		
+		index = self._index_by_name(name)
+		item = (hexsha, mode, name)
+		if index == -1:
+			self._cache.append(item)
+		else:
+			if force:
+				self._cache[index] = item
+			else:
+				ex_item = self._cache[index]
+				if ex_item[0] != hexsha or ex_item[1] != mode:
+					raise ValueError("Item %r existed with different properties" % name)
+				# END handle mismatch
+			# END handle force
+		# END handle name exists
+		return self
+		
+	def add_unchecked(self, hexsha, mode, name):
+		"""Add the given item to the tree, its correctness is assumed, which 
+		puts the caller into responsibility to assure the input is correct. 
+		For more information on the parameters, see ``add``"""
+		self._cache.append((hexsha, mode, name))
+		
+	def __delitem__(self, name):
+		"""Deletes an item with the given name if it exists"""
+		index = self._index_by_name(name)
+		if index > -1:
+			del(self._cache[index])
+		
+	#} END mutators
+
+
 class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializable):
 	"""
 	Tress represent a ordered list of Blobs and other Trees. Hence it can be 
@@ -42,8 +120,8 @@ class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializabl
 	type = "tree"
 	__slots__ = "_cache"
 	
-	# using ascii codes for comparison 
-	commit_id = 016		
+	# actual integer ids for comparison 
+	commit_id = 016
 	blob_id = 010
 	symlink_id = 012
 	tree_id = 004
@@ -56,7 +134,7 @@ class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializabl
 						}
 	
 	
-	def __init__(self, repo, sha, mode=0, path=None):
+	def __init__(self, repo, sha, mode=tree_id<<12, path=None):
 		super(Tree, self).__init__(repo, sha, mode, path)
 
 	@classmethod
@@ -133,7 +211,6 @@ class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializabl
 			yield (sha_to_hex(sha), mode, name)
 		# END for each byte in data stream
 
-
 	def __div__(self, file):
 		"""
 		Find the named object in this tree's contents
@@ -198,6 +275,13 @@ class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializabl
 		"""
 		return [ i for i in self if i.type == "blob" ]
 
+	@property
+	def cache(self):
+		""":return: An object allowing to modify the internal cache. This can be used
+		to change the tree's contents. When done, make sure you call ``set_done``
+		on the tree modifier, or serialization behaviour will be incorrect.
+		See the ``TreeModifier`` for more information on how to alter the cache"""
+		return TreeModifier(self._cache)
 
 	def traverse( self, predicate = lambda i,d: True,
 						   prune = lambda i,d: False, depth = -1, branch_first=True,
@@ -253,11 +337,9 @@ class Tree(base.IndexObject, diff.Diffable, utils.Traversable, utils.Serializabl
 		
 	def _serialize(self, stream, presort=False):
 		"""Serialize this tree into the stream. Please note that we will assume 
-		our tree data to be in a sorted state. If this is not the case, set the 
-		presort flag True
-		:param presort: if True, default False, sort our tree information before
-			writing it to the stream. This should be done if the cache changed
-			in the meanwhile"""
+		our tree data to be in a sorted state. If this is not the case, serialization
+		will not generate a correct tree representation as these are assumed to be sorted
+		by algorithms"""
 		ord_zero = ord('0')
 		bit_mask = 7			# 3 bits set
 		hex_to_bin = binascii.a2b_hex
diff --git a/test/git/test_tree.py b/test/git/test_tree.py
index 3f8ab851..d983cb2f 100644
--- a/test/git/test_tree.py
+++ b/test/git/test_tree.py
@@ -18,11 +18,12 @@ class TestTree(TestBase):
 			if item.type != Tree.type:
 				continue
 			# END skip non-trees
-			orig_data = item.data
-			orig_cache = item._cache
+			tree = item
+			orig_data = tree.data
+			orig_cache = tree._cache
 			
 			stream = StringIO()
-			item._serialize(stream)
+			tree._serialize(stream)
 			assert stream.getvalue() == orig_data
 			
 			stream.seek(0)
@@ -30,8 +31,60 @@ class TestTree(TestBase):
 			testtree._deserialize(stream)
 			assert testtree._cache == orig_cache
 			
-			# add an item, serialize with presort
-			self.fail("presort")
+			
+			# TEST CACHE MUTATOR
+			mod = testtree.cache
+			self.failUnlessRaises(ValueError, mod.add, "invalid sha", 0, "name")
+			self.failUnlessRaises(ValueError, mod.add, Tree.NULL_HEX_SHA, 0, "invalid mode")
+			self.failUnlessRaises(ValueError, mod.add, Tree.NULL_HEX_SHA, tree.mode, "invalid/name")
+			
+			# add new item
+			name = "fake_dir"
+			mod.add(testtree.NULL_HEX_SHA, tree.mode, name)
+			assert name in testtree
+			
+			# its available in the tree immediately
+			assert isinstance(testtree[name], Tree)
+			
+			# adding it again will not cause multiple of them to be presents
+			cur_count = len(testtree)
+			mod.add(testtree.NULL_HEX_SHA, tree.mode, name)
+			assert len(testtree) == cur_count
+			
+			# fails with a different sha - name exists
+			hexsha = "1"*40
+			self.failUnlessRaises(ValueError, mod.add, hexsha, tree.mode, name)
+			
+			# force it - replace existing one
+			mod.add(hexsha, tree.mode, name, force=True)
+			assert testtree[name].sha == hexsha
+			assert len(testtree) == cur_count
+			
+			# unchecked addition always works, even with invalid items
+			invalid_name = "hi/there"
+			mod.add_unchecked(hexsha, 0, invalid_name)
+			assert len(testtree) == cur_count + 1
+			
+			del(mod[invalid_name])
+			assert len(testtree) == cur_count
+			# del again, its fine
+			del(mod[invalid_name])
+			
+			# have added one item, we are done
+			mod.set_done()
+			mod.set_done()		# multiple times are okay
+			
+			# serialize, its different now
+			stream = StringIO()
+			testtree._serialize(stream)
+			stream.seek(0)
+			assert stream.getvalue() != orig_data
+			
+			# replaces cache, but we make sure of it
+			del(testtree._cache)
+			testtree._deserialize(stream)
+			assert name in testtree
+			assert invalid_name not in testtree
 		# END for each item in tree
 	
 	def test_traverse(self):