comparison mercurial/bundlerepo.py @ 12044:bcc7139521b7

bundlerepo: remove duplication of bundle decompressors
author Matt Mackall <mpm@selenic.com>
date Wed, 25 Aug 2010 16:55:54 -0500
parents 17031fea4e95
children 40935b59518b
comparison
equal deleted inserted replaced
12043:bef5effb3db0 12044:bcc7139521b7
11 were part of the actual repository. 11 were part of the actual repository.
12 """ 12 """
13 13
14 from node import nullid 14 from node import nullid
15 from i18n import _ 15 from i18n import _
16 import os, struct, bz2, zlib, tempfile, shutil 16 import os, struct, tempfile, shutil
17 import changegroup, util, mdiff 17 import changegroup, util, mdiff
18 import localrepo, changelog, manifest, filelog, revlog, error 18 import localrepo, changelog, manifest, filelog, revlog, error
19 19
20 class bundlerevlog(revlog.revlog): 20 class bundlerevlog(revlog.revlog):
21 def __init__(self, opener, indexfile, bundlefile, 21 def __init__(self, opener, indexfile, bundlefile,
170 else: 170 else:
171 self._url = 'bundle:' + bundlename 171 self._url = 'bundle:' + bundlename
172 172
173 self.tempfile = None 173 self.tempfile = None
174 self.bundlefile = open(bundlename, "rb") 174 self.bundlefile = open(bundlename, "rb")
175 header = self.bundlefile.read(6) 175 b = changegroup.readbundle(self.bundlefile, bundlename)
176 if not header.startswith("HG"): 176 if b.compressed():
177 raise util.Abort(_("%s: not a Mercurial bundle file") % bundlename)
178 elif not header.startswith("HG10"):
179 raise util.Abort(_("%s: unknown bundle version") % bundlename)
180 elif (header == "HG10BZ") or (header == "HG10GZ"):
181 fdtemp, temp = tempfile.mkstemp(prefix="hg-bundle-", 177 fdtemp, temp = tempfile.mkstemp(prefix="hg-bundle-",
182 suffix=".hg10un", dir=self.path) 178 suffix=".hg10un", dir=self.path)
183 self.tempfile = temp 179 self.tempfile = temp
184 fptemp = os.fdopen(fdtemp, 'wb') 180 fptemp = os.fdopen(fdtemp, 'wb')
185 def generator(f):
186 if header == "HG10BZ":
187 zd = bz2.BZ2Decompressor()
188 zd.decompress("BZ")
189 elif header == "HG10GZ":
190 zd = zlib.decompressobj()
191 for chunk in f:
192 yield zd.decompress(chunk)
193 gen = generator(util.filechunkiter(self.bundlefile, 4096))
194 181
195 try: 182 try:
196 fptemp.write("HG10UN") 183 fptemp.write("HG10UN")
197 for chunk in gen: 184 while 1:
185 chunk = b.read(2**18)
186 if not chunk:
187 break
198 fptemp.write(chunk) 188 fptemp.write(chunk)
199 finally: 189 finally:
200 fptemp.close() 190 fptemp.close()
201 self.bundlefile.close() 191 self.bundlefile.close()
202 192
203 self.bundlefile = open(self.tempfile, "rb") 193 self.bundlefile = open(self.tempfile, "rb")
204 # seek right after the header
205 self.bundlefile.seek(6) 194 self.bundlefile.seek(6)
206 elif header == "HG10UN": 195
207 # nothing to do
208 pass
209 else:
210 raise util.Abort(_("%s: unknown bundle compression type")
211 % bundlename)
212 # dict with the mapping 'filename' -> position in the bundle 196 # dict with the mapping 'filename' -> position in the bundle
213 self.bundlefilespos = {} 197 self.bundlefilespos = {}
214 198
215 @util.propertycache 199 @util.propertycache
216 def changelog(self): 200 def changelog(self):