diff mercurial/changegroup.py @ 3666:025f68f22ae2

move write_bundle to changegroup.py
author Matt Mackall <mpm@selenic.com>
date Wed, 15 Nov 2006 15:51:58 -0600
parents fe1689273f84
children 8500a13ec44b
line wrap: on
line diff
--- a/mercurial/changegroup.py
+++ b/mercurial/changegroup.py
@@ -8,7 +8,7 @@ of the GNU General Public License, incor
 """
 from i18n import gettext as _
 from demandload import *
-demandload(globals(), "struct util")
+demandload(globals(), "struct os bz2 util tempfile")
 
 def getchunk(source):
     """get a chunk from a changegroup"""
@@ -41,3 +41,55 @@ def genchunk(data):
 def closechunk():
     return struct.pack(">l", 0)
 
+class nocompress(object):
+    def compress(self, x):
+        return x
+    def flush(self):
+        return ""
+
+def writebundle(cg, filename, compress):
+    """Write a bundle file and return its filename.
+
+    Existing files will not be overwritten.
+    If no filename is specified, a temporary file is created.
+    bz2 compression can be turned off.
+    The bundle file will be deleted in case of errors.
+    """
+
+    fh = None
+    cleanup = None
+    try:
+        if filename:
+            if os.path.exists(filename):
+                raise util.Abort(_("file '%s' already exists") % filename)
+            fh = open(filename, "wb")
+        else:
+            fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")
+            fh = os.fdopen(fd, "wb")
+        cleanup = filename
+
+        if compress:
+            fh.write("HG10")
+            z = bz2.BZ2Compressor(9)
+        else:
+            fh.write("HG10UN")
+            z = nocompress()
+        # parse the changegroup data, otherwise we will block
+        # in case of sshrepo because we don't know the end of the stream
+
+        # an empty chunkiter is the end of the changegroup
+        empty = False
+        while not empty:
+            empty = True
+            for chunk in chunkiter(cg):
+                empty = False
+                fh.write(z.compress(genchunk(chunk)))
+            fh.write(z.compress(closechunk()))
+        fh.write(z.flush())
+        cleanup = None
+        return filename
+    finally:
+        if fh is not None:
+            fh.close()
+        if cleanup is not None:
+            os.unlink(cleanup)