1
# Copyright (C) 2005, 2006 Canonical Ltd
 
 
2
# Written by Robert Collins <robert.collins@canonical.com>
 
 
4
# This program is free software; you can redistribute it and/or modify
 
 
5
# it under the terms of the GNU General Public License as published by
 
 
6
# the Free Software Foundation; either version 2 of the License, or
 
 
7
# (at your option) any later version.
 
 
9
# This program is distributed in the hope that it will be useful,
 
 
10
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
 
11
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
 
12
# GNU General Public License for more details.
 
 
14
# You should have received a copy of the GNU General Public License
 
 
15
# along with this program; if not, write to the Free Software
 
 
16
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 
18
"""Bzrlib specific gzip tunings. We plan to feed these to the upstream gzip."""
 
 
20
from cStringIO import StringIO
 
 
22
# make GzipFile faster:
 
 
24
from gzip import FEXTRA, FCOMMENT, FNAME, FHCRC
 
 
29
# we want a \n preserved, break on \n only splitlines.
 
 
32
__all__ = ["GzipFile", "bytes_to_gzip"]
 
 
36
    """Return i as an unsigned integer, assuming it fits in 32 bits.
 
 
38
    If it's >= 2GB when viewed as a 32-bit unsigned int, return a long.
 
 
46
    """Return the low-order 32 bits of an int, as a non-negative int."""
 
 
47
    return i & 0xFFFFFFFFL
 
 
50
def bytes_to_gzip(bytes, factory=zlib.compressobj,
 
 
51
    level=zlib.Z_DEFAULT_COMPRESSION, method=zlib.DEFLATED,
 
 
52
    width=-zlib.MAX_WBITS, mem=zlib.DEF_MEM_LEVEL,
 
 
54
    """Create a gzip file containing bytes and return its content."""
 
 
56
        '\037\213'  # self.fileobj.write('\037\213')  # magic header
 
 
57
        '\010'      # self.fileobj.write('\010')      # compression method
 
 
58
                    # fname = self.filename[:-3]
 
 
62
        '\x00'      # self.fileobj.write(chr(flags))
 
 
63
        '\0\0\0\0'  # write32u(self.fileobj, long(time.time()))
 
 
64
        '\002'      # self.fileobj.write('\002')
 
 
65
        '\377'      # self.fileobj.write('\377')
 
 
67
        ''          #     self.fileobj.write(fname + '\000')
 
 
69
    # using a compressobj avoids a small header and trailer that the compress()
 
 
70
    # utility function adds.
 
 
71
    compress = factory(level, method, width, mem, 0)
 
 
72
    result.append(compress.compress(bytes))
 
 
73
    result.append(compress.flush())
 
 
74
    result.append(struct.pack("<L", LOWU32(crc32(bytes))))
 
 
75
    # size may exceed 2GB, or even 4GB
 
 
76
    result.append(struct.pack("<L", LOWU32(len(bytes))))
 
 
77
    return ''.join(result)
 
 
80
class GzipFile(gzip.GzipFile):
 
 
81
    """Knit tuned version of GzipFile.
 
 
83
    This is based on the following lsprof stats:
 
 
84
    python 2.4 stock GzipFile write:
 
 
85
    58971      0   5644.3090   2721.4730   gzip:193(write)
 
 
86
    +58971     0   1159.5530   1159.5530   +<built-in method compress>
 
 
87
    +176913    0    987.0320    987.0320   +<len>
 
 
88
    +58971     0    423.1450    423.1450   +<zlib.crc32>
 
 
89
    +58971     0    353.1060    353.1060   +<method 'write' of 'cStringIO.
 
 
92
    58971      0   4477.2590   2103.1120   bzrlib.knit:1250(write)
 
 
93
    +58971     0   1297.7620   1297.7620   +<built-in method compress>
 
 
94
    +58971     0    406.2160    406.2160   +<zlib.crc32>
 
 
95
    +58971     0    341.9020    341.9020   +<method 'write' of 'cStringIO.
 
 
97
    +58971     0    328.2670    328.2670   +<len>
 
 
100
    Yes, its only 1.6 seconds, but they add up.
 
 
103
    def _add_read_data(self, data):
 
 
105
        # temp var for len(data) and switch to +='s.
 
 
108
        self.crc = zlib.crc32(data, self.crc)
 
 
109
        self.extrabuf += data
 
 
110
        self.extrasize += len_data
 
 
111
        self.size += len_data
 
 
113
    def _write_gzip_header(self):
 
 
114
        """A tuned version of gzip._write_gzip_header
 
 
116
        We have some extra constrains that plain Gzip does not.
 
 
117
        1) We want to write the whole blob at once. rather than multiple 
 
 
118
           calls to fileobj.write().
 
 
119
        2) We never have a filename
 
 
120
        3) We don't care about the time
 
 
123
           '\037\213'   # self.fileobj.write('\037\213')  # magic header
 
 
124
            '\010'      # self.fileobj.write('\010')      # compression method
 
 
125
                        # fname = self.filename[:-3]
 
 
129
            '\x00'      # self.fileobj.write(chr(flags))
 
 
130
            '\0\0\0\0'  # write32u(self.fileobj, long(time.time()))
 
 
131
            '\002'      # self.fileobj.write('\002')
 
 
132
            '\377'      # self.fileobj.write('\377')
 
 
134
            ''          #     self.fileobj.write(fname + '\000')
 
 
137
    def _read(self, size=1024):
 
 
138
        # various optimisations:
 
 
139
        # reduces lsprof count from 2500 to 
 
 
140
        # 8337 calls in 1272, 365 internal
 
 
141
        if self.fileobj is None:
 
 
142
            raise EOFError, "Reached EOF"
 
 
145
            # If the _new_member flag is set, we have to
 
 
146
            # jump to the next member, if there is one.
 
 
148
            # First, check if we're at the end of the file;
 
 
149
            # if so, it's time to stop; no more members to read.
 
 
150
            next_header_bytes = self.fileobj.read(10)
 
 
151
            if next_header_bytes == '':
 
 
152
                raise EOFError, "Reached EOF"
 
 
155
            self._read_gzip_header(next_header_bytes)
 
 
156
            self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
 
 
157
            self._new_member = False
 
 
159
        # Read a chunk of data from the file
 
 
160
        buf = self.fileobj.read(size)
 
 
162
        # If the EOF has been reached, flush the decompression object
 
 
163
        # and mark this object as finished.
 
 
166
            self._add_read_data(self.decompress.flush())
 
 
167
            if len(self.decompress.unused_data) < 8:
 
 
168
                raise AssertionError("what does flush do?")
 
 
169
            self._gzip_tail = self.decompress.unused_data[0:8]
 
 
171
            # tell the driving read() call we have stuffed all the data
 
 
173
            raise EOFError, 'Reached EOF'
 
 
175
        self._add_read_data(self.decompress.decompress(buf))
 
 
177
        if self.decompress.unused_data != "":
 
 
178
            # Ending case: we've come to the end of a member in the file,
 
 
179
            # so seek back to the start of the data for the next member which
 
 
180
            # is the length of the decompress objects unused data - the first
 
 
181
            # 8 bytes for the end crc and size records.
 
 
183
            # so seek back to the start of the unused data, finish up
 
 
184
            # this member, and read a new gzip header.
 
 
185
            # (The number of bytes to seek back is the length of the unused
 
 
186
            # data, minus 8 because those 8 bytes are part of this member.
 
 
187
            seek_length = len (self.decompress.unused_data) - 8
 
 
189
                # we read too much data
 
 
190
                self.fileobj.seek(-seek_length, 1)
 
 
191
                self._gzip_tail = self.decompress.unused_data[0:8]
 
 
192
            elif seek_length < 0:
 
 
193
                # we haven't read enough to check the checksum.
 
 
194
                if not (-8 < seek_length):
 
 
195
                    raise AssertionError("too great a seek")
 
 
196
                buf = self.fileobj.read(-seek_length)
 
 
197
                self._gzip_tail = self.decompress.unused_data + buf
 
 
199
                self._gzip_tail = self.decompress.unused_data
 
 
201
            # Check the CRC and file size, and set the flag so we read
 
 
202
            # a new member on the next call
 
 
204
            self._new_member = True
 
 
207
        """tuned to reduce function calls and eliminate file seeking:
 
 
209
        reduces lsprof count from 800 to 288
 
 
211
        avoid U32 call by using struct format L
 
 
214
        # We've read to the end of the file, so we should have 8 bytes of 
 
 
215
        # unused data in the decompressor. If we don't, there is a corrupt file.
 
 
216
        # We use these 8 bytes to calculate the CRC and the recorded file size.
 
 
217
        # We then check the that the computed CRC and size of the
 
 
218
        # uncompressed data matches the stored values.  Note that the size
 
 
219
        # stored is the true file size mod 2**32.
 
 
220
        if not (len(self._gzip_tail) == 8):
 
 
221
            raise AssertionError("gzip trailer is incorrect length.")
 
 
222
        crc32, isize = struct.unpack("<LL", self._gzip_tail)
 
 
223
        # note that isize is unsigned - it can exceed 2GB
 
 
224
        if crc32 != U32(self.crc):
 
 
225
            raise IOError, "CRC check failed %d %d" % (crc32, U32(self.crc))
 
 
226
        elif isize != LOWU32(self.size):
 
 
227
            raise IOError, "Incorrect length of data produced"
 
 
229
    def _read_gzip_header(self, bytes=None):
 
 
230
        """Supply bytes if the minimum header size is already read.
 
 
232
        :param bytes: 10 bytes of header data.
 
 
234
        """starting cost: 300 in 3998
 
 
235
        15998 reads from 3998 calls
 
 
239
            bytes = self.fileobj.read(10)
 
 
241
        if magic != '\037\213':
 
 
242
            raise IOError, 'Not a gzipped file'
 
 
243
        method = ord(bytes[2:3])
 
 
245
            raise IOError, 'Unknown compression method'
 
 
246
        flag = ord(bytes[3:4])
 
 
247
        # modtime = self.fileobj.read(4) (bytes [4:8])
 
 
248
        # extraflag = self.fileobj.read(1) (bytes[8:9])
 
 
249
        # os = self.fileobj.read(1) (bytes[9:10])
 
 
250
        # self.fileobj.read(6)
 
 
253
            # Read & discard the extra field, if present
 
 
254
            xlen = ord(self.fileobj.read(1))
 
 
255
            xlen = xlen + 256*ord(self.fileobj.read(1))
 
 
256
            self.fileobj.read(xlen)
 
 
258
            # Read and discard a null-terminated string containing the filename
 
 
260
                s = self.fileobj.read(1)
 
 
261
                if not s or s=='\000':
 
 
264
            # Read and discard a null-terminated string containing a comment
 
 
266
                s = self.fileobj.read(1)
 
 
267
                if not s or s=='\000':
 
 
270
            self.fileobj.read(2)     # Read & discard the 16-bit header CRC
 
 
272
    def readline(self, size=-1):
 
 
273
        """Tuned to remove buffer length calls in _unread and...
 
 
275
        also removes multiple len(c) calls, inlines _unread,
 
 
276
        total savings - lsprof 5800 to 5300
 
 
279
        8176 calls to read() in 1684
 
 
280
        changing the min chunk size to 200 halved all the cache misses
 
 
281
        leading to a drop to:
 
 
283
        4168 call to read() in 1646
 
 
284
        - i.e. just reduced the function call overhead. May be worth 
 
 
287
        if size < 0: size = sys.maxint
 
 
289
        readsize = min(200, size)    # Read from the file in small chunks
 
 
292
                return "".join(bufs) # Return resulting line
 
 
295
            c = self.read(readsize)
 
 
296
            # number of bytes read
 
 
300
                # We set i=size to break out of the loop under two
 
 
301
                # conditions: 1) there's no newline, and the chunk is
 
 
302
                # larger than size, or 2) there is a newline, but the
 
 
303
                # resulting line would be longer than 'size'.
 
 
304
                if i==-1 and len_c > size: i=size-1
 
 
305
                elif size <= i: i = size -1
 
 
307
            if i >= 0 or c == '':
 
 
308
                # if i>= 0 we have a newline or have triggered the above
 
 
309
                # if size is not None condition.
 
 
310
                # if c == '' its EOF.
 
 
311
                bufs.append(c[:i+1])    # Add portion of last chunk
 
 
312
                # -- inlined self._unread --
 
 
313
                ## self._unread(c[i+1:], len_c - i)   # Push back rest of chunk
 
 
314
                self.extrabuf = c[i+1:] + self.extrabuf
 
 
315
                self.extrasize = len_c - i + self.extrasize
 
 
316
                self.offset -= len_c - i
 
 
317
                # -- end inlined self._unread --
 
 
318
                return ''.join(bufs)    # Return resulting line
 
 
320
            # Append chunk to list, decrease 'size',
 
 
323
            readsize = min(size, readsize * 2)
 
 
325
    def readlines(self, sizehint=0):
 
 
326
        # optimise to avoid all the buffer manipulation
 
 
327
        # lsprof changed from:
 
 
328
        # 4168 calls in 5472 with 32000 calls to readline()
 
 
331
        # Negative numbers result in reading all the lines
 
 
333
        # python's gzip routine uses sizehint. This is a more efficient way
 
 
334
        # than python uses to honor it. But it is even more efficient to
 
 
335
        # just read the entire thing and use cStringIO to split into lines.
 
 
338
        # content = self.read(sizehint)
 
 
339
        # return bzrlib.osutils.split_lines(content)
 
 
340
        content = StringIO(self.read(-1))
 
 
341
        return content.readlines()
 
 
343
    def _unread(self, buf, len_buf=None):
 
 
344
        """tuned to remove unneeded len calls.
 
 
346
        because this is such an inner routine in readline, and readline is
 
 
347
        in many inner loops, this has been inlined into readline().
 
 
349
        The len_buf parameter combined with the reduction in len calls dropped
 
 
350
        the lsprof ms count for this routine on my test data from 800 to 200 - 
 
 
355
        self.extrabuf = buf + self.extrabuf
 
 
356
        self.extrasize = len_buf + self.extrasize
 
 
357
        self.offset -= len_buf
 
 
359
    def write(self, data):
 
 
360
        if self.mode != gzip.WRITE:
 
 
362
            raise IOError(errno.EBADF, "write() on read-only GzipFile object")
 
 
364
        if self.fileobj is None:
 
 
365
            raise ValueError, "write() on closed GzipFile object"
 
 
368
            self.size = self.size + data_len
 
 
369
            self.crc = zlib.crc32(data, self.crc)
 
 
370
            self.fileobj.write( self.compress.compress(data) )
 
 
371
            self.offset += data_len
 
 
373
    def writelines(self, lines):
 
 
374
        # profiling indicated a significant overhead 
 
 
375
        # calling write for each line.
 
 
376
        # this batch call is a lot faster :).
 
 
377
        # (4 seconds to 1 seconds for the sample upgrades I was testing).
 
 
378
        self.write(''.join(lines))