1 | """Functions that read and write gzipped files. |
---|
2 | |
---|
3 | The user of the file doesn't have to worry about the compression, |
---|
4 | but random access is not allowed.""" |
---|
5 | |
---|
6 | # based on Andrew Kuchling's minigzip.py distributed with the zlib module |
---|
7 | |
---|
8 | import struct, sys, time, os |
---|
9 | import zlib |
---|
10 | import io |
---|
11 | import __builtin__ |
---|
12 | |
---|
13 | __all__ = ["GzipFile","open"] |
---|
14 | |
---|
15 | FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16 |
---|
16 | |
---|
17 | READ, WRITE = 1, 2 |
---|
18 | |
---|
19 | def write32u(output, value): |
---|
20 | # The L format writes the bit pattern correctly whether signed |
---|
21 | # or unsigned. |
---|
22 | output.write(struct.pack("<L", value)) |
---|
23 | |
---|
24 | def read32(input): |
---|
25 | return struct.unpack("<I", input.read(4))[0] |
---|
26 | |
---|
27 | def open(filename, mode="rb", compresslevel=9): |
---|
28 | """Shorthand for GzipFile(filename, mode, compresslevel). |
---|
29 | |
---|
30 | The filename argument is required; mode defaults to 'rb' |
---|
31 | and compresslevel defaults to 9. |
---|
32 | |
---|
33 | """ |
---|
34 | return GzipFile(filename, mode, compresslevel) |
---|
35 | |
---|
36 | class GzipFile(io.BufferedIOBase): |
---|
37 | """The GzipFile class simulates most of the methods of a file object with |
---|
38 | the exception of the readinto() and truncate() methods. |
---|
39 | |
---|
40 | """ |
---|
41 | |
---|
42 | myfileobj = None |
---|
43 | max_read_chunk = 10 * 1024 * 1024 # 10Mb |
---|
44 | |
---|
45 | def __init__(self, filename=None, mode=None, |
---|
46 | compresslevel=9, fileobj=None, mtime=None): |
---|
47 | """Constructor for the GzipFile class. |
---|
48 | |
---|
49 | At least one of fileobj and filename must be given a |
---|
50 | non-trivial value. |
---|
51 | |
---|
52 | The new class instance is based on fileobj, which can be a regular |
---|
53 | file, a StringIO object, or any other object which simulates a file. |
---|
54 | It defaults to None, in which case filename is opened to provide |
---|
55 | a file object. |
---|
56 | |
---|
57 | When fileobj is not None, the filename argument is only used to be |
---|
58 | included in the gzip file header, which may includes the original |
---|
59 | filename of the uncompressed file. It defaults to the filename of |
---|
60 | fileobj, if discernible; otherwise, it defaults to the empty string, |
---|
61 | and in this case the original filename is not included in the header. |
---|
62 | |
---|
63 | The mode argument can be any of 'r', 'rb', 'a', 'ab', 'w', or 'wb', |
---|
64 | depending on whether the file will be read or written. The default |
---|
65 | is the mode of fileobj if discernible; otherwise, the default is 'rb'. |
---|
66 | Be aware that only the 'rb', 'ab', and 'wb' values should be used |
---|
67 | for cross-platform portability. |
---|
68 | |
---|
69 | The compresslevel argument is an integer from 0 to 9 controlling the |
---|
70 | level of compression; 1 is fastest and produces the least compression, |
---|
71 | and 9 is slowest and produces the most compression. 0 is no compression |
---|
72 | at all. The default is 9. |
---|
73 | |
---|
74 | The mtime argument is an optional numeric timestamp to be written |
---|
75 | to the stream when compressing. All gzip compressed streams |
---|
76 | are required to contain a timestamp. If omitted or None, the |
---|
77 | current time is used. This module ignores the timestamp when |
---|
78 | decompressing; however, some programs, such as gunzip, make use |
---|
79 | of it. The format of the timestamp is the same as that of the |
---|
80 | return value of time.time() and of the st_mtime member of the |
---|
81 | object returned by os.stat(). |
---|
82 | |
---|
83 | """ |
---|
84 | |
---|
85 | # Make sure we don't inadvertently enable universal newlines on the |
---|
86 | # underlying file object - in read mode, this causes data corruption. |
---|
87 | if mode: |
---|
88 | mode = mode.replace('U', '') |
---|
89 | # guarantee the file is opened in binary mode on platforms |
---|
90 | # that care about that sort of thing |
---|
91 | if mode and 'b' not in mode: |
---|
92 | mode += 'b' |
---|
93 | if fileobj is None: |
---|
94 | fileobj = self.myfileobj = __builtin__.open(filename, mode or 'rb') |
---|
95 | if filename is None: |
---|
96 | # Issue #13781: os.fdopen() creates a fileobj with a bogus name |
---|
97 | # attribute. Avoid saving this in the gzip header's filename field. |
---|
98 | if hasattr(fileobj, 'name') and fileobj.name != '<fdopen>': |
---|
99 | filename = fileobj.name |
---|
100 | else: |
---|
101 | filename = '' |
---|
102 | if mode is None: |
---|
103 | if hasattr(fileobj, 'mode'): mode = fileobj.mode |
---|
104 | else: mode = 'rb' |
---|
105 | |
---|
106 | if mode[0:1] == 'r': |
---|
107 | self.mode = READ |
---|
108 | # Set flag indicating start of a new member |
---|
109 | self._new_member = True |
---|
110 | # Buffer data read from gzip file. extrastart is offset in |
---|
111 | # stream where buffer starts. extrasize is number of |
---|
112 | # bytes remaining in buffer from current stream position. |
---|
113 | self.extrabuf = "" |
---|
114 | self.extrasize = 0 |
---|
115 | self.extrastart = 0 |
---|
116 | self.name = filename |
---|
117 | # Starts small, scales exponentially |
---|
118 | self.min_readsize = 100 |
---|
119 | |
---|
120 | elif mode[0:1] == 'w' or mode[0:1] == 'a': |
---|
121 | self.mode = WRITE |
---|
122 | self._init_write(filename) |
---|
123 | self.compress = zlib.compressobj(compresslevel, |
---|
124 | zlib.DEFLATED, |
---|
125 | -zlib.MAX_WBITS, |
---|
126 | zlib.DEF_MEM_LEVEL, |
---|
127 | 0) |
---|
128 | else: |
---|
129 | raise IOError, "Mode " + mode + " not supported" |
---|
130 | |
---|
131 | self.fileobj = fileobj |
---|
132 | self.offset = 0 |
---|
133 | self.mtime = mtime |
---|
134 | |
---|
135 | if self.mode == WRITE: |
---|
136 | self._write_gzip_header() |
---|
137 | |
---|
138 | @property |
---|
139 | def filename(self): |
---|
140 | import warnings |
---|
141 | warnings.warn("use the name attribute", DeprecationWarning, 2) |
---|
142 | if self.mode == WRITE and self.name[-3:] != ".gz": |
---|
143 | return self.name + ".gz" |
---|
144 | return self.name |
---|
145 | |
---|
146 | def __repr__(self): |
---|
147 | s = repr(self.fileobj) |
---|
148 | return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>' |
---|
149 | |
---|
150 | def _check_closed(self): |
---|
151 | """Raises a ValueError if the underlying file object has been closed. |
---|
152 | |
---|
153 | """ |
---|
154 | if self.closed: |
---|
155 | raise ValueError('I/O operation on closed file.') |
---|
156 | |
---|
157 | def _init_write(self, filename): |
---|
158 | self.name = filename |
---|
159 | self.crc = zlib.crc32("") & 0xffffffffL |
---|
160 | self.size = 0 |
---|
161 | self.writebuf = [] |
---|
162 | self.bufsize = 0 |
---|
163 | |
---|
164 | def _write_gzip_header(self): |
---|
165 | self.fileobj.write('\037\213') # magic header |
---|
166 | self.fileobj.write('\010') # compression method |
---|
167 | try: |
---|
168 | # RFC 1952 requires the FNAME field to be Latin-1. Do not |
---|
169 | # include filenames that cannot be represented that way. |
---|
170 | fname = os.path.basename(self.name) |
---|
171 | if not isinstance(fname, str): |
---|
172 | fname = fname.encode('latin-1') |
---|
173 | if fname.endswith('.gz'): |
---|
174 | fname = fname[:-3] |
---|
175 | except UnicodeEncodeError: |
---|
176 | fname = '' |
---|
177 | flags = 0 |
---|
178 | if fname: |
---|
179 | flags = FNAME |
---|
180 | self.fileobj.write(chr(flags)) |
---|
181 | mtime = self.mtime |
---|
182 | if mtime is None: |
---|
183 | mtime = time.time() |
---|
184 | write32u(self.fileobj, long(mtime)) |
---|
185 | self.fileobj.write('\002') |
---|
186 | self.fileobj.write('\377') |
---|
187 | if fname: |
---|
188 | self.fileobj.write(fname + '\000') |
---|
189 | |
---|
190 | def _init_read(self): |
---|
191 | self.crc = zlib.crc32("") & 0xffffffffL |
---|
192 | self.size = 0 |
---|
193 | |
---|
194 | def _read_gzip_header(self): |
---|
195 | magic = self.fileobj.read(2) |
---|
196 | if magic != '\037\213': |
---|
197 | raise IOError, 'Not a gzipped file' |
---|
198 | method = ord( self.fileobj.read(1) ) |
---|
199 | if method != 8: |
---|
200 | raise IOError, 'Unknown compression method' |
---|
201 | flag = ord( self.fileobj.read(1) ) |
---|
202 | self.mtime = read32(self.fileobj) |
---|
203 | # extraflag = self.fileobj.read(1) |
---|
204 | # os = self.fileobj.read(1) |
---|
205 | self.fileobj.read(2) |
---|
206 | |
---|
207 | if flag & FEXTRA: |
---|
208 | # Read & discard the extra field, if present |
---|
209 | xlen = ord(self.fileobj.read(1)) |
---|
210 | xlen = xlen + 256*ord(self.fileobj.read(1)) |
---|
211 | self.fileobj.read(xlen) |
---|
212 | if flag & FNAME: |
---|
213 | # Read and discard a null-terminated string containing the filename |
---|
214 | while True: |
---|
215 | s = self.fileobj.read(1) |
---|
216 | if not s or s=='\000': |
---|
217 | break |
---|
218 | if flag & FCOMMENT: |
---|
219 | # Read and discard a null-terminated string containing a comment |
---|
220 | while True: |
---|
221 | s = self.fileobj.read(1) |
---|
222 | if not s or s=='\000': |
---|
223 | break |
---|
224 | if flag & FHCRC: |
---|
225 | self.fileobj.read(2) # Read & discard the 16-bit header CRC |
---|
226 | |
---|
227 | def write(self,data): |
---|
228 | self._check_closed() |
---|
229 | if self.mode != WRITE: |
---|
230 | import errno |
---|
231 | raise IOError(errno.EBADF, "write() on read-only GzipFile object") |
---|
232 | |
---|
233 | if self.fileobj is None: |
---|
234 | raise ValueError, "write() on closed GzipFile object" |
---|
235 | |
---|
236 | # Convert data type if called by io.BufferedWriter. |
---|
237 | if isinstance(data, memoryview): |
---|
238 | data = data.tobytes() |
---|
239 | |
---|
240 | if len(data) > 0: |
---|
241 | self.fileobj.write(self.compress.compress(data)) |
---|
242 | self.size += len(data) |
---|
243 | self.crc = zlib.crc32(data, self.crc) & 0xffffffffL |
---|
244 | self.offset += len(data) |
---|
245 | |
---|
246 | return len(data) |
---|
247 | |
---|
248 | def read(self, size=-1): |
---|
249 | self._check_closed() |
---|
250 | if self.mode != READ: |
---|
251 | import errno |
---|
252 | raise IOError(errno.EBADF, "read() on write-only GzipFile object") |
---|
253 | |
---|
254 | if self.extrasize <= 0 and self.fileobj is None: |
---|
255 | return '' |
---|
256 | |
---|
257 | readsize = 1024 |
---|
258 | if size < 0: # get the whole thing |
---|
259 | try: |
---|
260 | while True: |
---|
261 | self._read(readsize) |
---|
262 | readsize = min(self.max_read_chunk, readsize * 2) |
---|
263 | except EOFError: |
---|
264 | size = self.extrasize |
---|
265 | else: # just get some more of it |
---|
266 | try: |
---|
267 | while size > self.extrasize: |
---|
268 | self._read(readsize) |
---|
269 | readsize = min(self.max_read_chunk, readsize * 2) |
---|
270 | except EOFError: |
---|
271 | if size > self.extrasize: |
---|
272 | size = self.extrasize |
---|
273 | |
---|
274 | offset = self.offset - self.extrastart |
---|
275 | chunk = self.extrabuf[offset: offset + size] |
---|
276 | self.extrasize = self.extrasize - size |
---|
277 | |
---|
278 | self.offset += size |
---|
279 | return chunk |
---|
280 | |
---|
281 | def _unread(self, buf): |
---|
282 | self.extrasize = len(buf) + self.extrasize |
---|
283 | self.offset -= len(buf) |
---|
284 | |
---|
285 | def _read(self, size=1024): |
---|
286 | if self.fileobj is None: |
---|
287 | raise EOFError, "Reached EOF" |
---|
288 | |
---|
289 | if self._new_member: |
---|
290 | # If the _new_member flag is set, we have to |
---|
291 | # jump to the next member, if there is one. |
---|
292 | # |
---|
293 | # First, check if we're at the end of the file; |
---|
294 | # if so, it's time to stop; no more members to read. |
---|
295 | pos = self.fileobj.tell() # Save current position |
---|
296 | self.fileobj.seek(0, 2) # Seek to end of file |
---|
297 | if pos == self.fileobj.tell(): |
---|
298 | raise EOFError, "Reached EOF" |
---|
299 | else: |
---|
300 | self.fileobj.seek( pos ) # Return to original position |
---|
301 | |
---|
302 | self._init_read() |
---|
303 | self._read_gzip_header() |
---|
304 | self.decompress = zlib.decompressobj(-zlib.MAX_WBITS) |
---|
305 | self._new_member = False |
---|
306 | |
---|
307 | # Read a chunk of data from the file |
---|
308 | buf = self.fileobj.read(size) |
---|
309 | |
---|
310 | # If the EOF has been reached, flush the decompression object |
---|
311 | # and mark this object as finished. |
---|
312 | |
---|
313 | if buf == "": |
---|
314 | uncompress = self.decompress.flush() |
---|
315 | self._read_eof() |
---|
316 | self._add_read_data( uncompress ) |
---|
317 | raise EOFError, 'Reached EOF' |
---|
318 | |
---|
319 | uncompress = self.decompress.decompress(buf) |
---|
320 | self._add_read_data( uncompress ) |
---|
321 | |
---|
322 | if self.decompress.unused_data != "": |
---|
323 | # Ending case: we've come to the end of a member in the file, |
---|
324 | # so seek back to the start of the unused data, finish up |
---|
325 | # this member, and read a new gzip header. |
---|
326 | # (The number of bytes to seek back is the length of the unused |
---|
327 | # data, minus 8 because _read_eof() will rewind a further 8 bytes) |
---|
328 | self.fileobj.seek( -len(self.decompress.unused_data)+8, 1) |
---|
329 | |
---|
330 | # Check the CRC and file size, and set the flag so we read |
---|
331 | # a new member on the next call |
---|
332 | self._read_eof() |
---|
333 | self._new_member = True |
---|
334 | |
---|
335 | def _add_read_data(self, data): |
---|
336 | self.crc = zlib.crc32(data, self.crc) & 0xffffffffL |
---|
337 | offset = self.offset - self.extrastart |
---|
338 | self.extrabuf = self.extrabuf[offset:] + data |
---|
339 | self.extrasize = self.extrasize + len(data) |
---|
340 | self.extrastart = self.offset |
---|
341 | self.size = self.size + len(data) |
---|
342 | |
---|
343 | def _read_eof(self): |
---|
344 | # We've read to the end of the file, so we have to rewind in order |
---|
345 | # to reread the 8 bytes containing the CRC and the file size. |
---|
346 | # We check the that the computed CRC and size of the |
---|
347 | # uncompressed data matches the stored values. Note that the size |
---|
348 | # stored is the true file size mod 2**32. |
---|
349 | self.fileobj.seek(-8, 1) |
---|
350 | crc32 = read32(self.fileobj) |
---|
351 | isize = read32(self.fileobj) # may exceed 2GB |
---|
352 | if crc32 != self.crc: |
---|
353 | raise IOError("CRC check failed %s != %s" % (hex(crc32), |
---|
354 | hex(self.crc))) |
---|
355 | elif isize != (self.size & 0xffffffffL): |
---|
356 | raise IOError, "Incorrect length of data produced" |
---|
357 | |
---|
358 | # Gzip files can be padded with zeroes and still have archives. |
---|
359 | # Consume all zero bytes and set the file position to the first |
---|
360 | # non-zero byte. See http://www.gzip.org/#faq8 |
---|
361 | c = "\x00" |
---|
362 | while c == "\x00": |
---|
363 | c = self.fileobj.read(1) |
---|
364 | if c: |
---|
365 | self.fileobj.seek(-1, 1) |
---|
366 | |
---|
367 | @property |
---|
368 | def closed(self): |
---|
369 | return self.fileobj is None |
---|
370 | |
---|
371 | def close(self): |
---|
372 | fileobj = self.fileobj |
---|
373 | if fileobj is None: |
---|
374 | return |
---|
375 | self.fileobj = None |
---|
376 | try: |
---|
377 | if self.mode == WRITE: |
---|
378 | fileobj.write(self.compress.flush()) |
---|
379 | write32u(fileobj, self.crc) |
---|
380 | # self.size may exceed 2GB, or even 4GB |
---|
381 | write32u(fileobj, self.size & 0xffffffffL) |
---|
382 | finally: |
---|
383 | myfileobj = self.myfileobj |
---|
384 | if myfileobj: |
---|
385 | self.myfileobj = None |
---|
386 | myfileobj.close() |
---|
387 | |
---|
388 | def flush(self,zlib_mode=zlib.Z_SYNC_FLUSH): |
---|
389 | self._check_closed() |
---|
390 | if self.mode == WRITE: |
---|
391 | # Ensure the compressor's buffer is flushed |
---|
392 | self.fileobj.write(self.compress.flush(zlib_mode)) |
---|
393 | self.fileobj.flush() |
---|
394 | |
---|
395 | def fileno(self): |
---|
396 | """Invoke the underlying file object's fileno() method. |
---|
397 | |
---|
398 | This will raise AttributeError if the underlying file object |
---|
399 | doesn't support fileno(). |
---|
400 | """ |
---|
401 | return self.fileobj.fileno() |
---|
402 | |
---|
403 | def rewind(self): |
---|
404 | '''Return the uncompressed stream file position indicator to the |
---|
405 | beginning of the file''' |
---|
406 | if self.mode != READ: |
---|
407 | raise IOError("Can't rewind in write mode") |
---|
408 | self.fileobj.seek(0) |
---|
409 | self._new_member = True |
---|
410 | self.extrabuf = "" |
---|
411 | self.extrasize = 0 |
---|
412 | self.extrastart = 0 |
---|
413 | self.offset = 0 |
---|
414 | |
---|
415 | def readable(self): |
---|
416 | return self.mode == READ |
---|
417 | |
---|
418 | def writable(self): |
---|
419 | return self.mode == WRITE |
---|
420 | |
---|
421 | def seekable(self): |
---|
422 | return True |
---|
423 | |
---|
424 | def seek(self, offset, whence=0): |
---|
425 | if whence: |
---|
426 | if whence == 1: |
---|
427 | offset = self.offset + offset |
---|
428 | else: |
---|
429 | raise ValueError('Seek from end not supported') |
---|
430 | if self.mode == WRITE: |
---|
431 | if offset < self.offset: |
---|
432 | raise IOError('Negative seek in write mode') |
---|
433 | count = offset - self.offset |
---|
434 | for i in xrange(count // 1024): |
---|
435 | self.write(1024 * '\0') |
---|
436 | self.write((count % 1024) * '\0') |
---|
437 | elif self.mode == READ: |
---|
438 | if offset < self.offset: |
---|
439 | # for negative seek, rewind and do positive seek |
---|
440 | self.rewind() |
---|
441 | count = offset - self.offset |
---|
442 | for i in xrange(count // 1024): |
---|
443 | self.read(1024) |
---|
444 | self.read(count % 1024) |
---|
445 | |
---|
446 | return self.offset |
---|
447 | |
---|
448 | def readline(self, size=-1): |
---|
449 | if size < 0: |
---|
450 | # Shortcut common case - newline found in buffer. |
---|
451 | offset = self.offset - self.extrastart |
---|
452 | i = self.extrabuf.find('\n', offset) + 1 |
---|
453 | if i > 0: |
---|
454 | self.extrasize -= i - offset |
---|
455 | self.offset += i - offset |
---|
456 | return self.extrabuf[offset: i] |
---|
457 | |
---|
458 | size = sys.maxint |
---|
459 | readsize = self.min_readsize |
---|
460 | else: |
---|
461 | readsize = size |
---|
462 | bufs = [] |
---|
463 | while size != 0: |
---|
464 | c = self.read(readsize) |
---|
465 | i = c.find('\n') |
---|
466 | |
---|
467 | # We set i=size to break out of the loop under two |
---|
468 | # conditions: 1) there's no newline, and the chunk is |
---|
469 | # larger than size, or 2) there is a newline, but the |
---|
470 | # resulting line would be longer than 'size'. |
---|
471 | if (size <= i) or (i == -1 and len(c) > size): |
---|
472 | i = size - 1 |
---|
473 | |
---|
474 | if i >= 0 or c == '': |
---|
475 | bufs.append(c[:i + 1]) # Add portion of last chunk |
---|
476 | self._unread(c[i + 1:]) # Push back rest of chunk |
---|
477 | break |
---|
478 | |
---|
479 | # Append chunk to list, decrease 'size', |
---|
480 | bufs.append(c) |
---|
481 | size = size - len(c) |
---|
482 | readsize = min(size, readsize * 2) |
---|
483 | if readsize > self.min_readsize: |
---|
484 | self.min_readsize = min(readsize, self.min_readsize * 2, 512) |
---|
485 | return ''.join(bufs) # Return resulting line |
---|
486 | |
---|
487 | |
---|
488 | def _test(): |
---|
489 | # Act like gzip; with -d, act like gunzip. |
---|
490 | # The input file is not deleted, however, nor are any other gzip |
---|
491 | # options or features supported. |
---|
492 | args = sys.argv[1:] |
---|
493 | decompress = args and args[0] == "-d" |
---|
494 | if decompress: |
---|
495 | args = args[1:] |
---|
496 | if not args: |
---|
497 | args = ["-"] |
---|
498 | for arg in args: |
---|
499 | if decompress: |
---|
500 | if arg == "-": |
---|
501 | f = GzipFile(filename="", mode="rb", fileobj=sys.stdin) |
---|
502 | g = sys.stdout |
---|
503 | else: |
---|
504 | if arg[-3:] != ".gz": |
---|
505 | print "filename doesn't end in .gz:", repr(arg) |
---|
506 | continue |
---|
507 | f = open(arg, "rb") |
---|
508 | g = __builtin__.open(arg[:-3], "wb") |
---|
509 | else: |
---|
510 | if arg == "-": |
---|
511 | f = sys.stdin |
---|
512 | g = GzipFile(filename="", mode="wb", fileobj=sys.stdout) |
---|
513 | else: |
---|
514 | f = __builtin__.open(arg, "rb") |
---|
515 | g = open(arg + ".gz", "wb") |
---|
516 | while True: |
---|
517 | chunk = f.read(1024) |
---|
518 | if not chunk: |
---|
519 | break |
---|
520 | g.write(chunk) |
---|
521 | if g is not sys.stdout: |
---|
522 | g.close() |
---|
523 | if f is not sys.stdin: |
---|
524 | f.close() |
---|
525 | |
---|
526 | if __name__ == '__main__': |
---|
527 | _test() |
---|