|
28
|
1 # changegroup.py - Mercurial changegroup manipulation functions
|
|
|
2 #
|
|
|
3 # Copyright 2006 Matt Mackall <mpm@selenic.com>
|
|
|
4 #
|
|
|
5 # This software may be used and distributed according to the terms of the
|
|
|
6 # GNU General Public License version 2, incorporated herein by reference.
|
|
|
7
|
|
|
8 from i18n import _
|
|
|
9 import util
|
|
|
10 import struct, os, bz2, zlib, tempfile
|
|
|
11
|
|
|
12 def getchunk(source):
|
|
|
13 """get a chunk from a changegroup"""
|
|
|
14 d = source.read(4)
|
|
|
15 if not d:
|
|
|
16 return ""
|
|
|
17 l = struct.unpack(">l", d)[0]
|
|
|
18 if l <= 4:
|
|
|
19 return ""
|
|
|
20 d = source.read(l - 4)
|
|
|
21 if len(d) < l - 4:
|
|
|
22 raise util.Abort(_("premature EOF reading chunk"
|
|
|
23 " (got %d bytes, expected %d)")
|
|
|
24 % (len(d), l - 4))
|
|
|
25 return d
|
|
|
26
|
|
|
27 def chunkiter(source):
|
|
|
28 """iterate through the chunks in source"""
|
|
|
29 while 1:
|
|
|
30 c = getchunk(source)
|
|
|
31 if not c:
|
|
|
32 break
|
|
|
33 yield c
|
|
|
34
|
|
|
35 def chunkheader(length):
|
|
|
36 """build a changegroup chunk header"""
|
|
|
37 return struct.pack(">l", length + 4)
|
|
|
38
|
|
|
39 def closechunk():
|
|
|
40 return struct.pack(">l", 0)
|
|
|
41
|
|
|
42 class nocompress(object):
|
|
|
43 def compress(self, x):
|
|
|
44 return x
|
|
|
45 def flush(self):
|
|
|
46 return ""
|
|
|
47
|
|
|
48 bundletypes = {
|
|
|
49 "": ("", nocompress),
|
|
|
50 "HG10UN": ("HG10UN", nocompress),
|
|
|
51 "HG10BZ": ("HG10", lambda: bz2.BZ2Compressor()),
|
|
|
52 "HG10GZ": ("HG10GZ", lambda: zlib.compressobj()),
|
|
|
53 }
|
|
|
54
|
|
|
55 # hgweb uses this list to communicate it's preferred type
|
|
|
56 bundlepriority = ['HG10GZ', 'HG10BZ', 'HG10UN']
|
|
|
57
|
|
|
58 def writebundle(cg, filename, bundletype):
|
|
|
59 """Write a bundle file and return its filename.
|
|
|
60
|
|
|
61 Existing files will not be overwritten.
|
|
|
62 If no filename is specified, a temporary file is created.
|
|
|
63 bz2 compression can be turned off.
|
|
|
64 The bundle file will be deleted in case of errors.
|
|
|
65 """
|
|
|
66
|
|
|
67 fh = None
|
|
|
68 cleanup = None
|
|
|
69 try:
|
|
|
70 if filename:
|
|
|
71 fh = open(filename, "wb")
|
|
|
72 else:
|
|
|
73 fd, filename = tempfile.mkstemp(prefix="hg-bundle-", suffix=".hg")
|
|
|
74 fh = os.fdopen(fd, "wb")
|
|
|
75 cleanup = filename
|
|
|
76
|
|
|
77 header, compressor = bundletypes[bundletype]
|
|
|
78 fh.write(header)
|
|
|
79 z = compressor()
|
|
|
80
|
|
|
81 # parse the changegroup data, otherwise we will block
|
|
|
82 # in case of sshrepo because we don't know the end of the stream
|
|
|
83
|
|
|
84 # an empty chunkiter is the end of the changegroup
|
|
|
85 # a changegroup has at least 2 chunkiters (changelog and manifest).
|
|
|
86 # after that, an empty chunkiter is the end of the changegroup
|
|
|
87 empty = False
|
|
|
88 count = 0
|
|
|
89 while not empty or count <= 2:
|
|
|
90 empty = True
|
|
|
91 count += 1
|
|
|
92 for chunk in chunkiter(cg):
|
|
|
93 empty = False
|
|
|
94 fh.write(z.compress(chunkheader(len(chunk))))
|
|
|
95 pos = 0
|
|
|
96 while pos < len(chunk):
|
|
|
97 next = pos + 2**20
|
|
|
98 fh.write(z.compress(chunk[pos:next]))
|
|
|
99 pos = next
|
|
|
100 fh.write(z.compress(closechunk()))
|
|
|
101 fh.write(z.flush())
|
|
|
102 cleanup = None
|
|
|
103 return filename
|
|
|
104 finally:
|
|
|
105 if fh is not None:
|
|
|
106 fh.close()
|
|
|
107 if cleanup is not None:
|
|
|
108 os.unlink(cleanup)
|
|
|
109
|
|
|
110 def unbundle(header, fh):
|
|
|
111 if header == 'HG10UN':
|
|
|
112 return fh
|
|
|
113 elif not header.startswith('HG'):
|
|
|
114 # old client with uncompressed bundle
|
|
|
115 def generator(f):
|
|
|
116 yield header
|
|
|
117 for chunk in f:
|
|
|
118 yield chunk
|
|
|
119 elif header == 'HG10GZ':
|
|
|
120 def generator(f):
|
|
|
121 zd = zlib.decompressobj()
|
|
|
122 for chunk in f:
|
|
|
123 yield zd.decompress(chunk)
|
|
|
124 elif header == 'HG10BZ':
|
|
|
125 def generator(f):
|
|
|
126 zd = bz2.BZ2Decompressor()
|
|
|
127 zd.decompress("BZ")
|
|
|
128 for chunk in util.filechunkiter(f, 4096):
|
|
|
129 yield zd.decompress(chunk)
|
|
|
130 return util.chunkbuffer(generator(fh))
|
|
|
131
|
|
|
132 def readbundle(fh, fname):
|
|
|
133 header = fh.read(6)
|
|
|
134 if not header.startswith('HG'):
|
|
|
135 raise util.Abort(_('%s: not a Mercurial bundle file') % fname)
|
|
|
136 if not header.startswith('HG10'):
|
|
|
137 raise util.Abort(_('%s: unknown bundle version') % fname)
|
|
|
138 elif header not in bundletypes:
|
|
|
139 raise util.Abort(_('%s: unknown bundle compression type') % fname)
|
|
|
140 return unbundle(header, fh)
|
