create a readbundle function

mercurial/changegroup.py

@@ -93,3 +93,23 b' def writebundle(cg, filename, compress):'
             fh.close()
         if cleanup is not None:
             os.unlink(cleanup)
+
+def readbundle(fh):
+    header = fh.read(6)
+    if not header.startswith("HG"):
+        raise util.Abort(_("%s: not a Mercurial bundle file") % fname)
+    elif not header.startswith("HG10"):
+        raise util.Abort(_("%s: unknown bundle version") % fname)
+
+    if header == "HG10BZ":
+        def generator(f):
+            zd = bz2.BZ2Decompressor()
+            zd.decompress("BZ")
+            for chunk in util.filechunkiter(f, 4096):
+                yield zd.decompress(chunk)
+        return util.chunkbuffer(generator(fh))
+    elif header == "HG10UN":
+        return fh
+
+    raise util.Abort(_("%s: unknown bundle compression type")
+                     % fname)

mercurial/commands.py

@@ -11,7 +11,7 b' from i18n import gettext as _'
 demandload(globals(), "os re sys signal imp urllib pdb shlex")
 demandload(globals(), "fancyopts ui hg util lock revlog bundlerepo")
 demandload(globals(), "difflib patch time")
-demandload(globals(), "traceback errno version atexit bz2")
+demandload(globals(), "traceback errno version atexit")
 demandload(globals(), "archival changegroup cmdutil hgweb.server sshserver")
 
 class UnknownCommand(Exception):
@@ -2195,29 +2195,8 b' def unbundle(ui, repo, fname, **opts):'
     Apply a compressed changegroup file generated by the bundle
     command.
     """
-    f = urllib.urlopen(fname)
-
-    header = f.read(6)
-    if not header.startswith("HG"):
-        raise util.Abort(_("%s: not a Mercurial bundle file") % fname)
-    elif not header.startswith("HG10"):
-        raise util.Abort(_("%s: unknown bundle version") % fname)
-    elif header == "HG10BZ":
-        def generator(f):
-            zd = bz2.BZ2Decompressor()
-            zd.decompress("BZ")
-            for chunk in f:
-                yield zd.decompress(chunk)
-    elif header == "HG10UN":
-        def generator(f):
-            for chunk in f:
-                yield chunk
-    else:
-        raise util.Abort(_("%s: unknown bundle compression type")
-                         % fname)
-    gen = generator(util.filechunkiter(f, 4096))
-    modheads = repo.addchangegroup(util.chunkbuffer(gen), 'unbundle',
-                                   'bundle:' + fname)
+    gen = changegroup.readbundle(urllib.urlopen(fname))
+    modheads = repo.addchangegroup(gen, 'unbundle', 'bundle:' + fname)
     return postincoming(ui, repo, modheads, opts['update'])
 
 def update(ui, repo, node=None, merge=False, clean=False, force=None,