unduplicate bundle writing code from httprepo

mercurial/changegroup.py

@@ -8,7 +8,7 b' of the GNU General Public License, incor'
 """
 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 b' class nocompress(object):'
     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 b' def writebundle(cg, filename, compress):'
             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
 

mercurial/commands.py

@@ -332,7 +332,7 b' def bundle(ui, repo, fname, dest=None, *'
         cg = repo.changegroupsubset(o, revs, 'bundle')
     else:
         cg = repo.changegroup(o, 'bundle')
-    changegroup.writebundle(cg, fname, False)
+    changegroup.writebundle(cg, fname, "HG10")
 
 def cat(ui, repo, file1, *pats, **opts):
     """output the latest or given revisions of files
@@ -1292,7 +1292,8 b' def incoming(ui, repo, source="default",'
         if fname or not other.local():
             # create a bundle (uncompressed if other repo is not local)
             cg = other.changegroup(incoming, "incoming")
-            fname = cleanup = changegroup.writebundle(cg, fname, other.local())
+            type = other.local() and "HG10" or "HG10UN"
+            fname = cleanup = changegroup.writebundle(cg, fname, type)
             # keep written bundle?
             if opts["bundle"]:
                 cleanup = None

mercurial/httprepo.py

@@ -11,7 +11,7 b' from remoterepo import *'
 from i18n import gettext as _
 from demandload import *
 demandload(globals(), "hg os urllib urllib2 urlparse zlib util httplib")
-demandload(globals(), "errno keepalive tempfile socket")
+demandload(globals(), "errno keepalive tempfile socket changegroup")
 
 class passwordmgr(urllib2.HTTPPasswordMgrWithDefaultRealm):
     def __init__(self, ui):
@@ -326,47 +326,18 b' class httprepository(remoterepository):'
         # have to stream bundle to a temp file because we do not have
         # http 1.1 chunked transfer.
 
-        # XXX duplication from commands.py
-        class nocompress(object):
-            def compress(self, x):
-                return x
-            def flush(self):
-                return ""
-
-        unbundleversions = self.capable('unbundle')
-        try:
-            unbundleversions = unbundleversions.split(',')
-        except AttributeError:
-            unbundleversions = [""]
+        type = ""
+        types = self.capable('unbundle')
+        if types:
+            for x in types.split(','):
+                if x in changegroup.bundletypes:
+                    type = x
+                    break
 
-        while unbundleversions:
-            header = unbundleversions[0]
-            if header == "HG10GZ":
-                self.ui.note(_("using zlib compression\n"))
-                z = zlib.compressobj()
-                break
-            elif header == "HG10UN":
-                self.ui.note(_("using no compression\n"))
-                z = nocompress()
-                break
-            elif header == "":
-                self.ui.note(_("old server without compression support,"
-                               " sending uncompressed\n"))
-                z = nocompress()
-                break
-            unbundleversions.pop(0)
-        if not unbundleversions:
-            raise util.Abort(_("The server doesn't accept any bundle format"
-                               " method we know."))
-
-        fd, tempname = tempfile.mkstemp(prefix='hg-unbundle-')
-        fp = os.fdopen(fd, 'wb+')
+        tempname = changegroup.writebundle(cg, None, type)
+        fp = file(tempname, "rb")
         try:
-            fp.write(header)
-            for chunk in util.filechunkiter(cg):
-                fp.write(z.compress(chunk))
-            fp.write(z.flush())
-            length = fp.tell()
+            length = os.stat(tempname).st_size
             try:
                 rfp = self.do_cmd(
                     'unbundle', data=fp,