changegroup.py
153 lines
| 4.7 KiB
| text/x-python
|
PythonLexer
/ mercurial / changegroup.py
|
r8226 | # 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 | ||||
|
r10263 | # GNU General Public License version 2 or any later version. | ||
|
r3877 | |||
|
r3891 | from i18n import _ | ||
|
r8312 | import util | ||
import struct, os, bz2, zlib, tempfile | ||||
|
r1981 | |||
def getchunk(source): | ||||
|
r9437 | """return the next chunk from changegroup 'source' as a string""" | ||
|
r1981 | 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 | ||||
|
r10430 | def chunkiter(source, progress=None): | ||
|
r9437 | """iterate through the chunks in source, yielding a sequence of chunks | ||
(strings)""" | ||||
|
r1981 | while 1: | ||
c = getchunk(source) | ||||
if not c: | ||||
break | ||||
|
r10430 | elif progress is not None: | ||
progress() | ||||
|
r1981 | yield c | ||
|
r5368 | def chunkheader(length): | ||
|
r9437 | """return a changegroup chunk header (string)""" | ||
|
r5368 | return struct.pack(">l", length + 4) | ||
|
r1981 | |||
def closechunk(): | ||||
|
r9437 | """return a changegroup chunk header (string) for a zero-length chunk""" | ||
|
r1981 | return struct.pack(">l", 0) | ||
|
r3659 | class nocompress(object): | ||
def compress(self, x): | ||||
return x | ||||
def flush(self): | ||||
return "" | ||||
|
r3662 | bundletypes = { | ||
|
r3704 | "": ("", nocompress), | ||
"HG10UN": ("HG10UN", nocompress), | ||||
|
r3762 | "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()), | ||
"HG10GZ": ("HG10GZ", lambda: zlib.compressobj()), | ||||
|
r3662 | } | ||
|
r10356 | def collector(cl, mmfs, files): | ||
# Gather information about changeset nodes going out in a bundle. | ||||
# We want to gather manifests needed and filelogs affected. | ||||
def collect(node): | ||||
c = cl.read(node) | ||||
|
r11648 | files.update(c[3]) | ||
|
r10356 | mmfs.setdefault(c[0], node) | ||
return collect | ||||
|
r9087 | # hgweb uses this list to communicate its preferred type | ||
|
r6152 | bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN'] | ||
|
r3706 | def writebundle(cg, filename, bundletype): | ||
|
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: | ||||
fh = open(filename, "wb") | ||||
else: | ||||
fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg") | ||||
fh = os.fdopen(fd, "wb") | ||||
cleanup = filename | ||||
|
r3706 | header, compressor = bundletypes[bundletype] | ||
|
r3704 | fh.write(header) | ||
z = compressor() | ||||
|
r3662 | |||
|
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 | ||||
|
r5906 | # a changegroup has at least 2 chunkiters (changelog and manifest). | ||
# after that, an empty chunkiter is the end of the changegroup | ||||
|
r3659 | empty = False | ||
|
r5906 | count = 0 | ||
while not empty or count <= 2: | ||||
|
r3659 | empty = True | ||
|
r5906 | count += 1 | ||
|
r3659 | for chunk in chunkiter(cg): | ||
empty = False | ||||
|
r5368 | fh.write(z.compress(chunkheader(len(chunk)))) | ||
pos = 0 | ||||
while pos < len(chunk): | ||||
next = pos + 2**20 | ||||
fh.write(z.compress(chunk[pos:next])) | ||||
pos = next | ||||
|
r3659 | 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 | |||
|
r6154 | def unbundle(header, fh): | ||
if header == 'HG10UN': | ||||
return fh | ||||
elif not header.startswith('HG'): | ||||
# old client with uncompressed bundle | ||||
def generator(f): | ||||
yield header | ||||
for chunk in f: | ||||
yield chunk | ||||
elif header == 'HG10GZ': | ||||
def generator(f): | ||||
zd = zlib.decompressobj() | ||||
for chunk in f: | ||||
yield zd.decompress(chunk) | ||||
elif header == 'HG10BZ': | ||||
|
r3660 | def generator(f): | ||
zd = bz2.BZ2Decompressor() | ||||
zd.decompress("BZ") | ||||
for chunk in util.filechunkiter(f, 4096): | ||||
yield zd.decompress(chunk) | ||||
|
r6154 | return util.chunkbuffer(generator(fh)) | ||
|
r3660 | |||
|
r6154 | def readbundle(fh, fname): | ||
header = fh.read(6) | ||||
if not header.startswith('HG'): | ||||
raise util.Abort(_('%s: not a Mercurial bundle file') % fname) | ||||
if not header.startswith('HG10'): | ||||
raise util.Abort(_('%s: unknown bundle version') % fname) | ||||
elif header not in bundletypes: | ||||
raise util.Abort(_('%s: unknown bundle compression type') % fname) | ||||
return unbundle(header, fh) | ||||