changegroup.py
256 lines
| 8.1 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 _ | ||
|
r14143 | from node import nullrev | ||
import mdiff, util | ||||
|
r8312 | import struct, os, bz2, zlib, tempfile | ||
|
r1981 | |||
|
r14141 | _BUNDLE10_DELTA_HEADER = "20s20s20s20s" | ||
|
r13457 | def readexactly(stream, n): | ||
'''read n bytes from stream.read and abort if less was available''' | ||||
s = stream.read(n) | ||||
if len(s) < n: | ||||
raise util.Abort(_("stream ended unexpectedly" | ||||
" (got %d bytes, expected %d)") | ||||
% (len(s), n)) | ||||
return s | ||||
def getchunk(stream): | ||||
"""return the next chunk from stream as a string""" | ||||
d = readexactly(stream, 4) | ||||
|
r1981 | l = struct.unpack(">l", d)[0] | ||
if l <= 4: | ||||
|
r13458 | if l: | ||
raise util.Abort(_("invalid chunk length %d") % l) | ||||
|
r1981 | return "" | ||
|
r13457 | return readexactly(stream, l - 4) | ||
|
r1981 | |||
|
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 = { | ||
|
r14060 | "": ("", nocompress), # only when using unbundle on ssh and old http servers | ||
# since the unification ssh accepts a header but there | ||||
# is no capability signaling it. | ||||
|
r3704 | "HG10UN": ("HG10UN", nocompress), | ||
|
r3762 | "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()), | ||
"HG10GZ": ("HG10GZ", lambda: zlib.compressobj()), | ||||
|
r3662 | } | ||
|
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 | ||||
|
r12335 | # an empty chunkgroup is the end of the changegroup | ||
# a changegroup has at least 2 chunkgroups (changelog and manifest). | ||||
# after that, an empty chunkgroup is the end of the changegroup | ||||
|
r3659 | empty = False | ||
|
r5906 | count = 0 | ||
while not empty or count <= 2: | ||||
|
r3659 | empty = True | ||
|
r5906 | count += 1 | ||
|
r12335 | while 1: | ||
chunk = getchunk(cg) | ||||
if not chunk: | ||||
break | ||||
|
r3659 | 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 | |||
|
r12041 | def decompressor(fh, alg): | ||
if alg == 'UN': | ||||
|
r6154 | return fh | ||
|
r12041 | elif alg == 'GZ': | ||
|
r6154 | def generator(f): | ||
zd = zlib.decompressobj() | ||||
for chunk in f: | ||||
yield zd.decompress(chunk) | ||||
|
r12041 | elif alg == 'BZ': | ||
|
r3660 | def generator(f): | ||
zd = bz2.BZ2Decompressor() | ||||
zd.decompress("BZ") | ||||
for chunk in util.filechunkiter(f, 4096): | ||||
yield zd.decompress(chunk) | ||||
|
r12041 | else: | ||
raise util.Abort("unknown bundle compression '%s'" % alg) | ||||
|
r12329 | return util.chunkbuffer(generator(fh)) | ||
|
r12041 | |||
|
r12043 | class unbundle10(object): | ||
|
r14141 | deltaheader = _BUNDLE10_DELTA_HEADER | ||
deltaheadersize = struct.calcsize(deltaheader) | ||||
|
r12043 | def __init__(self, fh, alg): | ||
|
r12329 | self._stream = decompressor(fh, alg) | ||
|
r12044 | self._type = alg | ||
|
r12334 | self.callback = None | ||
|
r12044 | def compressed(self): | ||
return self._type != 'UN' | ||||
|
r12043 | def read(self, l): | ||
return self._stream.read(l) | ||||
|
r12330 | def seek(self, pos): | ||
return self._stream.seek(pos) | ||||
def tell(self): | ||||
|
r12332 | return self._stream.tell() | ||
|
r12347 | def close(self): | ||
return self._stream.close() | ||||
|
r12334 | |||
def chunklength(self): | ||||
|
r13459 | d = readexactly(self._stream, 4) | ||
|
r13458 | l = struct.unpack(">l", d)[0] | ||
if l <= 4: | ||||
if l: | ||||
raise util.Abort(_("invalid chunk length %d") % l) | ||||
return 0 | ||||
if self.callback: | ||||
|
r12334 | self.callback() | ||
|
r13458 | return l - 4 | ||
|
r12334 | |||
|
r14144 | def changelogheader(self): | ||
"""v10 does not have a changelog header chunk""" | ||||
return {} | ||||
def manifestheader(self): | ||||
"""v10 does not have a manifest header chunk""" | ||||
return {} | ||||
def filelogheader(self): | ||||
"""return the header of the filelogs chunk, v10 only has the filename""" | ||||
|
r12334 | l = self.chunklength() | ||
|
r14144 | if not l: | ||
return {} | ||||
fname = readexactly(self._stream, l) | ||||
return dict(filename=fname) | ||||
|
r12334 | |||
|
r14141 | def _deltaheader(self, headertuple, prevnode): | ||
node, p1, p2, cs = headertuple | ||||
if prevnode is None: | ||||
deltabase = p1 | ||||
else: | ||||
deltabase = prevnode | ||||
return node, p1, p2, deltabase, cs | ||||
|
r14144 | def deltachunk(self, prevnode): | ||
|
r12336 | l = self.chunklength() | ||
if not l: | ||||
return {} | ||||
|
r14141 | headerdata = readexactly(self._stream, self.deltaheadersize) | ||
header = struct.unpack(self.deltaheader, headerdata) | ||||
delta = readexactly(self._stream, l - self.deltaheadersize) | ||||
node, p1, p2, deltabase, cs = self._deltaheader(header, prevnode) | ||||
return dict(node=node, p1=p1, p2=p2, cs=cs, | ||||
deltabase=deltabase, delta=delta) | ||||
|
r12336 | |||
|
r12329 | class headerlessfixup(object): | ||
def __init__(self, fh, h): | ||||
self._h = h | ||||
self._fh = fh | ||||
def read(self, n): | ||||
if self._h: | ||||
d, self._h = self._h[:n], self._h[n:] | ||||
if len(d) < n: | ||||
|
r13457 | d += readexactly(self._fh, n - len(d)) | ||
|
r12329 | return d | ||
|
r13457 | return readexactly(self._fh, n) | ||
|
r12329 | |||
|
r6154 | def readbundle(fh, fname): | ||
|
r13457 | header = readexactly(fh, 6) | ||
|
r12042 | |||
if not fname: | ||||
fname = "stream" | ||||
if not header.startswith('HG') and header.startswith('\0'): | ||||
|
r12329 | fh = headerlessfixup(fh, header) | ||
|
r12042 | header = "HG10UN" | ||
magic, version, alg = header[0:2], header[2:4], header[4:6] | ||||
if magic != 'HG': | ||||
raise util.Abort(_('%s: not a Mercurial bundle') % fname) | ||||
if version != '10': | ||||
raise util.Abort(_('%s: unknown bundle version %s') % (fname, version)) | ||||
|
r12043 | return unbundle10(fh, alg) | ||
|
r13831 | |||
class bundle10(object): | ||||
|
r14143 | deltaheader = _BUNDLE10_DELTA_HEADER | ||
|
r13831 | def __init__(self, lookup): | ||
self._lookup = lookup | ||||
def close(self): | ||||
return closechunk() | ||||
def fileheader(self, fname): | ||||
return chunkheader(len(fname)) + fname | ||||
|
r14143 | def revchunk(self, revlog, rev, prev): | ||
node = revlog.node(rev) | ||||
p1, p2 = revlog.parentrevs(rev) | ||||
base = prev | ||||
prefix = '' | ||||
if base == nullrev: | ||||
delta = revlog.revision(node) | ||||
prefix = mdiff.trivialdiffheader(len(delta)) | ||||
else: | ||||
delta = revlog.revdiff(base, rev) | ||||
|
r13831 | linknode = self._lookup(revlog, node) | ||
|
r14143 | p1n, p2n = revlog.parents(node) | ||
basenode = revlog.node(base) | ||||
meta = self.builddeltaheader(node, p1n, p2n, basenode, linknode) | ||||
meta += prefix | ||||
l = len(meta) + len(delta) | ||||
|
r13831 | yield chunkheader(l) | ||
yield meta | ||||
|
r14143 | yield delta | ||
def builddeltaheader(self, node, p1n, p2n, basenode, linknode): | ||||
# do nothing with basenode, it is implicitly the previous one in HG10 | ||||
return struct.pack(self.deltaheader, node, p1n, p2n, linknode) | ||||