Improved archive function by allowing it to directly write to an output stream - previously it would cache everything to memory and try to provide zipping functionality itself

gitcmd: allows the output stream to be set explicitly which is mainly useful for archiving operations

2 files changed, 48 insertions, 55 deletions

diff --git a/lib/git/cmd.py b/lib/git/cmd.py
index ef7a9c6c..500fcd93 100644
--- a/lib/git/cmd.py
+++ b/lib/git/cmd.py
@@ -13,7 +13,8 @@ from errors import GitCommandError
 GIT_PYTHON_TRACE = os.environ.get("GIT_PYTHON_TRACE", False)
 
 execute_kwargs = ('istream', 'with_keep_cwd', 'with_extended_output',
-				  'with_exceptions', 'with_raw_output', 'as_process')
+				  'with_exceptions', 'with_raw_output', 'as_process', 
+				  'output_stream' )
 
 extra = {}
 if sys.platform == 'win32':
@@ -102,7 +103,8 @@ class Git(object):
 				with_extended_output=False,
 				with_exceptions=True,
 				with_raw_output=False,
-				as_process=False
+				as_process=False, 
+				output_stream=None
 				):
 		"""
 		Handles executing the command on the shell and consumes and returns
@@ -130,16 +132,20 @@ class Git(object):
 		``with_raw_output``
 			Whether to avoid stripping off trailing whitespace.
 			
-		 ``as_process``
-		 	Whether to return the created process instance directly from which 
-		 	streams can be read on demand. This will render with_extended_output, 
-		 	with_exceptions and with_raw_output ineffective - the caller will have 
-		 	to deal with the details himself.
-		 	It is important to note that the process will be placed into an AutoInterrupt
-		 	wrapper that will interrupt the process once it goes out of scope. If you 
-		 	use the command in iterators, you should pass the whole process instance 
-		 	instead of a single stream.
-
+		``as_process``
+			Whether to return the created process instance directly from which 
+			streams can be read on demand. This will render with_extended_output, 
+			with_exceptions and with_raw_output ineffective - the caller will have 
+			to deal with the details himself.
+			It is important to note that the process will be placed into an AutoInterrupt
+			wrapper that will interrupt the process once it goes out of scope. If you 
+			use the command in iterators, you should pass the whole process instance 
+			instead of a single stream.
+		``output_stream``
+			If set to a file-like object, data produced by the git command will be 
+			output to the given stream directly. 
+			Otherwise a new file will be opened.
+		
 		Returns::
 		
 		 str(output)								  # extended_output = False (Default)
@@ -160,13 +166,17 @@ class Git(object):
 		  cwd = os.getcwd()
 		else:
 		  cwd=self.git_dir
+		  
+		ostream = subprocess.PIPE
+		if output_stream is not None:
+			ostream = output_stream
 
 		# Start the process
 		proc = subprocess.Popen(command,
 								cwd=cwd,
 								stdin=istream,
 								stderr=subprocess.PIPE,
-								stdout=subprocess.PIPE,
+								stdout=ostream,
 								**extra
 								)
 
diff --git a/lib/git/repo.py b/lib/git/repo.py
index 39b1cb50..554c10cb 100644
--- a/lib/git/repo.py
+++ b/lib/git/repo.py
@@ -456,48 +456,27 @@ class Repo(object):
 		self.git.clone(self.path, path, **kwargs)
 		return Repo(path)
 
-	def archive_tar(self, treeish='master', prefix=None):
-		"""
-		Archive the given treeish
-
-		``treeish``
-			is the treeish name/id (default 'master')
-
-		``prefix``
-			is the optional prefix to prepend to each filename in the archive
-
-		Examples::
-
-			>>> repo.archive_tar
-			<String containing tar archive>
-
-			>>> repo.archive_tar('a87ff14')
-			<String containing tar archive for commit a87ff14>
 
-			>>> repo.archive_tar('master', 'myproject/')
-			<String containing tar bytes archive, whose files are prefixed with 'myproject/'>
-
-		Returns
-			str (containing bytes of tar archive)
+	def archive(self, ostream, treeish=None, prefix=None,  **kwargs):
 		"""
-		options = {}
-		if prefix:
-			options['prefix'] = prefix
-		return self.git.archive(treeish, **options)
-
-	def archive_tar_gz(self, treeish='master', prefix=None):
-		"""
-		Archive and gzip the given treeish
+		Archive the tree at the given revision.
+		``ostream``
+			file compatible stream object to which the archive will be written
 
 		``treeish``
-			is the treeish name/id (default 'master')
+			is the treeish name/id, defaults to active branch
 
 		``prefix``
 			is the optional prefix to prepend to each filename in the archive
+			
+		``kwargs``
+			Additional arguments passed to git-archive
+			NOTE: Use the 'format' argument to define the kind of format. Use 
+			specialized ostreams to write any format supported by python
 
 		Examples::
 
-			>>> repo.archive_tar_gz
+			>>> repo.archive(open("archive"
 			<String containing tar.gz archive>
 
 			>>> repo.archive_tar_gz('a87ff14')
@@ -506,18 +485,22 @@ class Repo(object):
 			>>> repo.archive_tar_gz('master', 'myproject/')
 			<String containing tar.gz archive and prefixed with 'myproject/'>
 
-		Returns
-			str (containing the bytes of tar.gz archive)
+		Raise
+			GitCommandError in case something went wrong
+			
 		"""
-		kwargs = {}
-		if prefix:
+		if treeish is None:
+			treeish = self.active_branch
+		if prefix and 'prefix' not in kwargs:
 			kwargs['prefix'] = prefix
-		resultstr =  self.git.archive(treeish, **kwargs)
-		sio = StringIO.StringIO()
-		gf = gzip.GzipFile(fileobj=sio, mode ='wb')
-		gf.write(resultstr)
-		gf.close()
-		return sio.getvalue()
+		kwargs['as_process'] = True
+		kwargs['output_stream'] = ostream
+		
+		proc =  self.git.archive(treeish, **kwargs)
+		status = proc.wait()
+		if status != 0:
+			raise GitCommandError( "git-archive", status, proc.stderr.read() )
+		
 
 
 	def __repr__(self):