index.add: Finished implemenation including through tests

index.checkout: added simple method allowing to checkout files from the index, including simple test

3 files changed, 224 insertions, 5 deletions

diff --git a/TODO b/TODO
index 69fbc307..d841f774 100644
--- a/TODO
+++ b/TODO
@@ -72,6 +72,8 @@ Index
   to keep the internal entry cache and write once everything is done. Problem 
   would be that all other git commands are unaware of the changes unless the index
   gets written. Its worth an evaluation at least.
+  A problem going with it is that there might be shell-related limitations on non-unix
+  where the commandline grows too large.
 * index.remove: On windows, there can be a command line length overflow 
   as we pass the paths directly as argv. This is as we use git-rm to be able 
   to remove whole directories easily. This could be implemented using 
diff --git a/lib/git/index.py b/lib/git/index.py
index 89e716d4..cc3f3a4e 100644
--- a/lib/git/index.py
+++ b/lib/git/index.py
@@ -15,6 +15,7 @@ import tempfile
 import os
 import sys
 import stat
+import subprocess
 import git.diff as diff
 
 from git.objects import Blob, Tree, Object, Commit
@@ -51,6 +52,9 @@ class BaseIndexEntry(tuple):
 	use numeric indices for performance reasons.
 	"""
 	
+	def __str__(self):
+		return "%o %s %i\t%s\n" % (self.mode, self.sha, self.stage, self.path)
+	
 	@property
 	def mode(self):
 		"""
@@ -466,6 +470,12 @@ class IndexFile(LazyMixin, diff.Diffable):
 		ret |= (index_mode & 0111)
 		return ret
 	
+	@classmethod
+	def _tree_mode_to_index_mode(cls, tree_mode):
+		"""
+		Convert a tree mode to index mode as good as possible
+		"""
+	
 	def iter_blobs(self, predicate = lambda t: True):
 		"""
 		Returns
@@ -597,9 +607,29 @@ class IndexFile(LazyMixin, diff.Diffable):
 			raise ValueError("Absolute path %r is not in git repository at %r" % (path,self.repo.git.git_dir))
 		return relative_path
 	
+	def _preprocess_add_items(self, items):
+		"""
+		Split the items into two lists of path strings and BaseEntries.
+		"""
+		paths = list()
+		entries = list()
+		
+		for item in items:
+			if isinstance(item, basestring):
+				paths.append(self._to_relative_path(item))
+			elif isinstance(item, Blob):
+				entries.append(BaseIndexEntry.from_blob(item))
+			elif isinstance(item, BaseIndexEntry):
+				entries.append(item)
+			else:
+				raise TypeError("Invalid Type: %r" % item)
+		# END for each item
+		return (paths, entries)
+		
+	
 	@clear_cache
 	@default_index
-	def add(self, items, **kwargs):
+	def add(self, items, force=True, **kwargs):
 		"""
 		Add files from the working tree, specific blobs or BaseIndexEntries 
 		to the index. The underlying index file will be written immediately, hence
@@ -612,18 +642,23 @@ class IndexFile(LazyMixin, diff.Diffable):
 			
 			- path string
 				strings denote a relative or absolute path into the repository pointing to 
-				an existing file, i.e. CHANGES, lib/myfile.ext, /home/gitrepo/lib/myfile.ext.
+				an existing file, i.e. CHANGES, lib/myfile.ext, '/home/gitrepo/lib/myfile.ext'.
 				
 				Paths provided like this must exist. When added, they will be written 
 				into the object database.
 				
+				PathStrings may contain globs, such as 'lib/__init__*' or can be directories
+				like 'lib', the latter ones will add all the files within the dirctory and 
+				subdirectories.
+				
 				This equals a straight git-add.
 				
 				They are added at stage 0
 				
 			- Blob object
 				Blobs are added as they are assuming a valid mode is set.
-				The file they refer to may or may not exist in the file system
+				The file they refer to may or may not exist in the file system, but 
+				must be a path relative to our repository.
 				
 				If their sha is null ( 40*0 ), their path must exist in the file system 
 				as an object will be created from the data at the path.The handling 
@@ -634,12 +669,21 @@ class IndexFile(LazyMixin, diff.Diffable):
 				is not dereferenced automatically, except that it can be created on 
 				filesystems not supporting it as well.
 				
+				Please note that globs or directories are not allowed in Blob objects. 
+				
 				They are added at stage 0
 				
 			- BaseIndexEntry or type
 				Handling equals the one of Blob objects, but the stage may be 
 				explicitly set.
 			
+		``force``
+			If True, otherwise ignored or excluded files will be 
+			added anyway.
+			As opposed to the git-add command, we enable this flag by default 
+			as the API user usually wants the item to be added even though 
+			they might be excluded.
+		
 		``**kwargs``
 			Additional keyword arguments to be passed to git-update-index, such 
 			as index_only.
@@ -647,7 +691,54 @@ class IndexFile(LazyMixin, diff.Diffable):
 		Returns
 			List(BaseIndexEntries) representing the entries just actually added.
 		"""
-		raise NotImplementedError("todo")
+		# sort the entries into strings and Entries, Blobs are converted to entries
+		# automatically
+		# paths can be git-added, for everything else we use git-update-index
+		entries_added = list()
+		paths, entries = self._preprocess_add_items(items)
+		
+		if paths:
+			git_add_output = self.repo.git.add(paths, v=True)
+			# force rereading our entries
+			del(self.entries)
+			for line in git_add_output.splitlines():
+				# line contains:
+				# add '<path>'
+				added_file = line[5:-1]
+				entries_added.append(self.entries[(added_file,0)])
+			# END for each line
+		# END path handling
+		
+		if entries:
+			null_mode_entries = [ e for e in entries if e.mode == 0 ]
+			if null_mode_entries:
+				raise ValueError("At least one Entry has a null-mode - please use index.remove to remove files for clarity")
+			# END null mode should be remove
+			
+			# create objects if required, otherwise go with the existing shas
+			null_entries_indices = [ i for i,e in enumerate(entries) if e.sha == Object.NULL_HEX_SHA ]
+			if null_entries_indices:
+				hash_proc = self.repo.git.hash_object(w=True, stdin_paths=True, istream=subprocess.PIPE, as_process=True)
+				hash_proc.stdin.write('\n'.join(entries[i].path for i in null_entries_indices))
+				obj_ids = self._flush_stdin_and_wait(hash_proc).splitlines()
+				assert len(obj_ids) == len(null_entries_indices), "git-hash-object did not produce all requested objects: want %i, got %i" % ( len(null_entries_indices), len(obj_ids) )
+				
+				# update IndexEntries with new object id
+				for i,new_sha in zip(null_entries_indices, obj_ids):
+					e = entries[i]
+					new_entry = BaseIndexEntry((e.mode, new_sha, e.stage, e.path))
+					entries[i] = new_entry
+				# END for each index
+			# END null_entry handling
+				
+			# feed all the data to stdin
+			update_index_proc = self.repo.git.update_index(index_info=True, istream=subprocess.PIPE, as_process=True, **kwargs)
+			update_index_proc.stdin.write('\n'.join(str(e) for e in entries))
+			entries_added.extend(entries)
+			self._flush_stdin_and_wait(update_index_proc)
+		# END if there are base entries
+		
+		return entries_added
 		
 	@clear_cache
 	@default_index
@@ -768,6 +859,53 @@ class IndexFile(LazyMixin, diff.Diffable):
 			fp.close()
 			os.remove(tmp_file_path)
 			
+	@classmethod
+	def _flush_stdin_and_wait(cls, proc):
+		proc.stdin.flush()
+		proc.stdin.close()
+		stdout = proc.stdout.read()
+		proc.wait()
+		return stdout
+	
+	@default_index
+	def checkout(self, paths=None, force=False, **kwargs):
+		"""
+		Checkout the given paths or all files from the version in the index.
+		
+		``paths``
+			If None, all paths in the index will be checked out. Otherwise an iterable
+			or single path of relative or absolute paths pointing to files is expected.
+			The command will ignore paths that do not exist.
+			
+		``force``
+			If True, existing files will be overwritten. If False, these will 
+			be skipped.
+			
+		``**kwargs``
+			Additional arguments to be pasesd to git-checkout-index
+			
+		Returns
+			self
+		"""
+		args = ["--index"]
+		if force:
+			args.append("--force")
+		
+		if paths is None:
+			args.append("--all")
+			self.repo.git.checkout_index(*args, **kwargs)
+		else:
+			if not isinstance(paths, (tuple,list)):
+				paths = [paths]
+				
+			args.append("--stdin")
+			paths = [self._to_relative_path(p) for p in paths]
+			co_proc = self.repo.git.checkout_index(args, as_process=True, istream=subprocess.PIPE, **kwargs)
+			co_proc.stdin.write('\n'.join(paths))
+			self._flush_stdin_and_wait(co_proc)
+		# END paths handling 
+		return self
+			
 	@clear_cache
 	@default_index
 	def reset(self, commit='HEAD', working_tree=False, paths=None, head=False, **kwargs):
diff --git a/test/git/test_index.py b/test/git/test_index.py
index e25f3d2e..3312abe1 100644
--- a/test/git/test_index.py
+++ b/test/git/test_index.py
@@ -8,8 +8,10 @@ from test.testlib import *
 from git import *
 import inspect
 import os
+import sys
 import tempfile
 import glob
+from stat import *
 
 class TestTree(TestBase):
 	
@@ -174,6 +176,26 @@ class TestTree(TestBase):
 			assert fp.read() != new_data
 		finally:
 			fp.close()
+			
+		# test full checkout
+		test_file = os.path.join(rw_repo.git.git_dir, "CHANGES")
+		os.remove(test_file)
+		index.checkout(None, force=True)
+		assert os.path.isfile(test_file)
+		
+		os.remove(test_file)
+		index.checkout(None, force=False)
+		assert os.path.isfile(test_file)
+		
+		# individual file
+		os.remove(test_file)
+		index.checkout(test_file)
+		assert os.path.exists(test_file)
+		
+		
+		
+		# currently it ignore non-existing paths
+		index.checkout(paths=["doesnt/exist"])
 		
 	
 	def _count_existing(self, repo, files):
@@ -186,6 +208,17 @@ class TestTree(TestBase):
 	# END num existing helper
 	
 		
+	def _make_file(self, rela_path, data, repo=None):
+		"""
+		Create a file at the given path relative to our repository, filled
+		with the given data. Returns absolute path to created file.
+		"""
+		repo = repo or self.rorepo
+		abs_path = os.path.join(repo.git.git_dir, rela_path)
+		fp = open(abs_path, "w")
+		fp.write(data)
+		fp.close()
+		return abs_path
 		
 	@with_rw_repo('0.1.6')
 	def test_index_mutation(self, rw_repo):
@@ -272,5 +305,51 @@ class TestTree(TestBase):
 		assert (lib_file_path, 0) not in index.entries
 		assert os.path.isfile(os.path.join(rw_repo.git.git_dir, lib_file_path))
 		
-		self.fail( "add file using simple path, blob, blob as symlink, entries with stages" )
+		# directory
+		entries = index.add(['lib'])
+		assert len(entries)>1
+		
+		# glob 
+		entries = index.reset(new_commit).add(['lib/*.py'])
+		assert len(entries) == 14
+		
+		# missing path
+		self.failUnlessRaises(GitCommandError, index.reset(new_commit).add, ['doesnt/exist/must/raise'])
+		
+		# blob from older revision overrides current index revision
+		old_blob = new_commit.parents[0].tree.blobs[0]
+		entries = index.reset(new_commit).add([old_blob])
+		assert index.entries[(old_blob.path,0)].sha == old_blob.id and len(entries) == 1 
+		
+		# mode 0 not allowed
+		null_sha = "0"*40
+		self.failUnlessRaises(ValueError, index.reset(new_commit).add, [BaseIndexEntry((0, null_sha,0,"doesntmatter"))])
+		
+		# add new file
+		new_file_relapath = "my_new_file"
+		new_file_path = self._make_file(new_file_relapath, "hello world", rw_repo)
+		entries = index.reset(new_commit).add([BaseIndexEntry((010644, null_sha, 0, new_file_relapath))])
+		assert len(entries) == 1 and entries[0].sha != null_sha
+		
+		# add symlink
+		if sys.platform != "win32":
+			link_file = os.path.join(rw_repo.git.git_dir, "my_real_symlink")
+			os.symlink("/etc/that", link_file)
+			entries = index.reset(new_commit).add([link_file])
+			assert len(entries) == 1 and S_ISLNK(entries[0].mode)
+			print "%o" % entries[0].mode
+		# END real symlink test 
+		
+		# add fake symlink and assure it checks-our as symlink
+		fake_symlink_relapath = "my_fake_symlink"
+		fake_symlink_path = self._make_file(fake_symlink_relapath, "/etc/that", rw_repo)
+		fake_entry = BaseIndexEntry((0120000, null_sha, 0, fake_symlink_relapath))
+		entries = index.reset(new_commit).add([fake_entry])
+		assert len(entries) == 1 and S_ISLNK(entries[0].mode)
+		
+		# checkout the fakelink, should be a link then
+		assert not S_ISLNK(os.stat(fake_symlink_path)[ST_MODE])
+		os.remove(fake_symlink_path)
+		index.checkout(fake_symlink_path)
+		assert S_ISLNK(os.lstat(fake_symlink_path)[ST_MODE])