changegroup.py
119 lines
| 3.3 KiB
| text/x-python
|
PythonLexer
/ mercurial / changegroup.py
|
r1981 | """ | ||
changegroup.py - Mercurial changegroup manipulation functions | ||||
Copyright 2006 Matt Mackall <mpm@selenic.com> | ||||
This software may be used and distributed according to the terms | ||||
of the GNU General Public License, incorporated herein by reference. | ||||
""" | ||||
|
r1997 | from i18n import gettext as _ | ||
|
r1981 | from demandload import * | ||
|
r3662 | demandload(globals(), "struct os bz2 zlib util tempfile") | ||
|
r1981 | |||
def getchunk(source): | ||||
"""get a chunk from a changegroup""" | ||||
d = source.read(4) | ||||
if not d: | ||||
return "" | ||||
l = struct.unpack(">l", d)[0] | ||||
if l <= 4: | ||||
return "" | ||||
d = source.read(l - 4) | ||||
if len(d) < l - 4: | ||||
raise util.Abort(_("premature EOF reading chunk" | ||||
" (got %d bytes, expected %d)") | ||||
% (len(d), l - 4)) | ||||
return d | ||||
def chunkiter(source): | ||||
"""iterate through the chunks in source""" | ||||
while 1: | ||||
c = getchunk(source) | ||||
if not c: | ||||
break | ||||
yield c | ||||
def genchunk(data): | ||||
"""build a changegroup chunk""" | ||||
header = struct.pack(">l", len(data)+ 4) | ||||
return "%s%s" % (header, data) | ||||
def closechunk(): | ||||
return struct.pack(">l", 0) | ||||
|
r3659 | class nocompress(object): | ||
def compress(self, x): | ||||
return x | ||||
def flush(self): | ||||
return "" | ||||
|
r3662 | bundletypes = { | ||
"": nocompress, | ||||
"HG10UN": nocompress, | ||||
"HG10": lambda: bz2.BZ2Compressor(9), | ||||
"HG10GZ": zlib.compressobj, | ||||
} | ||||
def writebundle(cg, filename, type): | ||||
|
r3659 | """Write a bundle file and return its filename. | ||
Existing files will not be overwritten. | ||||
If no filename is specified, a temporary file is created. | ||||
bz2 compression can be turned off. | ||||
The bundle file will be deleted in case of errors. | ||||
""" | ||||
fh = None | ||||
cleanup = None | ||||
try: | ||||
if filename: | ||||
if os.path.exists(filename): | ||||
raise util.Abort(_("file '%s' already exists") % filename) | ||||
fh = open(filename, "wb") | ||||
else: | ||||
fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg") | ||||
fh = os.fdopen(fd, "wb") | ||||
cleanup = filename | ||||
|
r3662 | fh.write(type) | ||
z = bundletypes[type]() | ||||
|
r3659 | # parse the changegroup data, otherwise we will block | ||
# in case of sshrepo because we don't know the end of the stream | ||||
# an empty chunkiter is the end of the changegroup | ||||
empty = False | ||||
while not empty: | ||||
empty = True | ||||
for chunk in chunkiter(cg): | ||||
empty = False | ||||
fh.write(z.compress(genchunk(chunk))) | ||||
fh.write(z.compress(closechunk())) | ||||
fh.write(z.flush()) | ||||
cleanup = None | ||||
return filename | ||||
finally: | ||||
if fh is not None: | ||||
fh.close() | ||||
if cleanup is not None: | ||||
os.unlink(cleanup) | ||||
|
r3660 | |||
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) | ||||