diff mercurial/changegroup.py @ 3662:f4dc02d7fb71

unduplicate bundle writing code from httprepo
author Matt Mackall <mpm@selenic.com>
date Wed, 15 Nov 2006 23:37:45 -0600
parents 8500a13ec44b
children 9c1737a3e254
line wrap: on
line diff
--- a/mercurial/changegroup.py	Wed Nov 15 15:51:58 2006 -0600
+++ b/mercurial/changegroup.py	Wed Nov 15 23:37:45 2006 -0600
@@ -8,7 +8,7 @@
 """
 from i18n import gettext as _
 from demandload import *
-demandload(globals(), "struct os bz2 util tempfile")
+demandload(globals(), "struct os bz2 zlib util tempfile")
 
 def getchunk(source):
     """get a chunk from a changegroup"""
@@ -47,7 +47,14 @@
     def flush(self):
         return ""
 
-def writebundle(cg, filename, compress):
+bundletypes = {
+    "": nocompress,
+    "HG10UN": nocompress,
+    "HG10": lambda: bz2.BZ2Compressor(9),
+    "HG10GZ": zlib.compressobj,
+}
+
+def writebundle(cg, filename, type):
     """Write a bundle file and return its filename.
 
     Existing files will not be overwritten.
@@ -68,12 +75,9 @@
             fh = os.fdopen(fd, "wb")
         cleanup = filename
 
-        if compress:
-            fh.write("HG10")
-            z = bz2.BZ2Compressor(9)
-        else:
-            fh.write("HG10UN")
-            z = nocompress()
+        fh.write(type)
+        z = bundletypes[type]()
+
         # parse the changegroup data, otherwise we will block
         # in case of sshrepo because we don't know the end of the stream