# HG changeset patch # User Matt Mackall # Date 1282773354 18000 # Node ID bcc7139521b77e6e8ac851a41ef1fca7769aaaa5 # Parent bef5effb3db0cc5b1366bc9c793e89a26a743125 bundlerepo: remove duplication of bundle decompressors diff -r bef5effb3db0 -r bcc7139521b7 mercurial/bundlerepo.py --- a/mercurial/bundlerepo.py Wed Aug 25 16:53:06 2010 -0500 +++ b/mercurial/bundlerepo.py Wed Aug 25 16:55:54 2010 -0500 @@ -13,7 +13,7 @@ from node import nullid from i18n import _ -import os, struct, bz2, zlib, tempfile, shutil +import os, struct, tempfile, shutil import changegroup, util, mdiff import localrepo, changelog, manifest, filelog, revlog, error @@ -172,43 +172,27 @@ self.tempfile = None self.bundlefile = open(bundlename, "rb") - header = self.bundlefile.read(6) - if not header.startswith("HG"): - raise util.Abort(_("%s: not a Mercurial bundle file") % bundlename) - elif not header.startswith("HG10"): - raise util.Abort(_("%s: unknown bundle version") % bundlename) - elif (header == "HG10BZ") or (header == "HG10GZ"): + b = changegroup.readbundle(self.bundlefile, bundlename) + if b.compressed(): fdtemp, temp = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg10un", dir=self.path) self.tempfile = temp fptemp = os.fdopen(fdtemp, 'wb') - def generator(f): - if header == "HG10BZ": - zd = bz2.BZ2Decompressor() - zd.decompress("BZ") - elif header == "HG10GZ": - zd = zlib.decompressobj() - for chunk in f: - yield zd.decompress(chunk) - gen = generator(util.filechunkiter(self.bundlefile, 4096)) try: fptemp.write("HG10UN") - for chunk in gen: + while 1: + chunk = b.read(2**18) + if not chunk: + break fptemp.write(chunk) finally: fptemp.close() self.bundlefile.close() self.bundlefile = open(self.tempfile, "rb") - # seek right after the header self.bundlefile.seek(6) - elif header == "HG10UN": - # nothing to do - pass - else: - raise util.Abort(_("%s: unknown bundle compression type") - % bundlename) + # dict with the mapping 'filename' -> position in the bundle self.bundlefilespos = {} diff -r bef5effb3db0 -r bcc7139521b7 mercurial/changegroup.py --- a/mercurial/changegroup.py Wed Aug 25 16:53:06 2010 -0500 +++ b/mercurial/changegroup.py Wed Aug 25 16:55:54 2010 -0500 @@ -141,6 +141,9 @@ class unbundle10(object): def __init__(self, fh, alg): self._stream = util.chunkbuffer(decompressor(fh, alg)) + self._type = alg + def compressed(self): + return self._type != 'UN' def read(self, l): return self._stream.read(l) diff -r bef5effb3db0 -r bcc7139521b7 tests/test-bundle-type.t --- a/tests/test-bundle-type.t Wed Aug 25 16:53:06 2010 -0500 +++ b/tests/test-bundle-type.t Wed Aug 25 16:55:54 2010 -0500 @@ -87,7 +87,7 @@ $ hg init tgarbage $ cd tgarbage $ hg pull ../bgarbage - abort: ../bgarbage: not a Mercurial bundle file + abort: ../bgarbage: not a Mercurial bundle $ cd .. test invalid bundle type