Mercurial > hg > python

comparison gzip_plus.py @ 76:cba9cb3a6797

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
from python3.13 with block boundary tracking added by HST
author Henry S Thompson <ht@inf.ed.ac.uk>
date Thu, 12 Mar 2026 18:06:10 +0000
parents
children
comparison
equal deleted inserted replaced
69:157f012ffab7 76:cba9cb3a6797
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 _compression
9 import builtins
10 import io
11 import os
12 import struct
13 import sys
14 import time
15 import weakref
16 import zlib
17
18 __all__ = ["BadGzipFile", "GzipFile", "open", "compress", "decompress"]
19
20 FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
21
22 READ = 'rb'
23 WRITE = 'wb'
24
25 _COMPRESS_LEVEL_FAST = 1
26 _COMPRESS_LEVEL_TRADEOFF = 6
27 _COMPRESS_LEVEL_BEST = 9
28
29 READ_BUFFER_SIZE = 128 * 1024
30 _WRITE_BUFFER_SIZE = 4 * io.DEFAULT_BUFFER_SIZE
31
32
33 def open(filename, mode="rb", compresslevel=_COMPRESS_LEVEL_BEST,
34 encoding=None, errors=None, newline=None):
35 """Open a gzip-compressed file in binary or text mode.
36
37 The filename argument can be an actual filename (a str or bytes object), or
38 an existing file object to read from or write to.
39
40 The mode argument can be "r", "rb", "w", "wb", "x", "xb", "a" or "ab" for
41 binary mode, or "rt", "wt", "xt" or "at" for text mode. The default mode is
42 "rb", and the default compresslevel is 9.
43
44 For binary mode, this function is equivalent to the GzipFile constructor:
45 GzipFile(filename, mode, compresslevel). In this case, the encoding, errors
46 and newline arguments must not be provided.
47
48 For text mode, a GzipFile object is created, and wrapped in an
49 io.TextIOWrapper instance with the specified encoding, error handling
50 behavior, and line ending(s).
51
52 """
53 if "t" in mode:
54 if "b" in mode:
55 raise ValueError("Invalid mode: %r" % (mode,))
56 else:
57 if encoding is not None:
58 raise ValueError("Argument 'encoding' not supported in binary mode")
59 if errors is not None:
60 raise ValueError("Argument 'errors' not supported in binary mode")
61 if newline is not None:
62 raise ValueError("Argument 'newline' not supported in binary mode")
63
64 gz_mode = mode.replace("t", "")
65 if isinstance(filename, (str, bytes, os.PathLike)):
66 binary_file = GzipFile(filename, gz_mode, compresslevel)
67 elif hasattr(filename, "read") or hasattr(filename, "write"):
68 binary_file = GzipFile(None, gz_mode, compresslevel, filename)
69 else:
70 raise TypeError("filename must be a str or bytes object, or a file")
71
72 if "t" in mode:
73 encoding = io.text_encoding(encoding)
74 return io.TextIOWrapper(binary_file, encoding, errors, newline)
75 else:
76 return binary_file
77
78 def write32u(output, value):
79 # The L format writes the bit pattern correctly whether signed
80 # or unsigned.
81 output.write(struct.pack("<L", value))
82
83 class _PaddedFile:
84 """Minimal read-only file object that prepends a string to the contents
85 of an actual file. Shouldn't be used outside of gzip.py, as it lacks
86 essential functionality."""
87
88 def __init__(self, f, prepend=b''):
89 #print('i',len(prepend))
90 self._buffer = prepend
91 self._length = len(prepend)
92 self.file = f
93 self._read = 0
94 self._pos = 0
95
96 def tell(self):
97 if self._read is None:
98 return self.file.tell()
99 else:
100 return self._pos
101
102 def read(self, size):
103 #print('r',size, self._read, self._length, self._pos)
104 if self._read is None:
105 res = self.file.read(size)
106 self._pos += len(res)
107 #print('r1', self._pos)
108 return res
109 if self._read + size <= self._length:
110 read = self._read
111 self._read += size
112 self._pos += size
113 #print('r2', self._pos)
114 return self._buffer[read:self._read]
115 else:
116 read = self._read
117 self._read = None
118 avail = self._length - read
119 balance = size - avail
120 res = self.file.read(balance)
121 self._pos += avail + len(res)
122 #print('r3', self._pos)
123 return self._buffer[read:] + res
124
125 def prepend(self, prepend=b''):
126 #print('p',len(prepend))
127 self._pos -= len(prepend)
128 if self._read is None:
129 self._buffer = prepend
130 else: # Assume data was read since the last prepend() call
131 self._read -= len(prepend)
132 return
133 self._length = len(self._buffer)
134 self._read = 0
135
136 def seek(self, off):
137 #print('o',off)
138 self._read = None
139 self._buffer = None
140 self._pos = self.file.seek(off)
141 return self._pos
142
143 def seekable(self):
144 return True # Allows fast-forwarding even in unseekable streams
145
146
147 class BadGzipFile(OSError):
148 """Exception raised in some cases for invalid gzip files."""
149
150
151 class _WriteBufferStream(io.RawIOBase):
152 """Minimal object to pass WriteBuffer flushes into GzipFile"""
153 def __init__(self, gzip_file):
154 self.gzip_file = weakref.ref(gzip_file)
155
156 def write(self, data):
157 gzip_file = self.gzip_file()
158 if gzip_file is None:
159 raise RuntimeError("lost gzip_file")
160 return gzip_file._write_raw(data)
161
162 def seekable(self):
163 return False
164
165 def writable(self):
166 return True
167
168
169 class GzipFile(_compression.BaseStream):
170 """The GzipFile class simulates most of the methods of a file object with
171 the exception of the truncate() method.
172
173 This class only supports opening files in binary mode. If you need to open a
174 compressed file in text mode, use the gzip.open() function.
175
176 """
177
178 # Overridden with internal file object to be closed, if only a filename
179 # is passed in
180 myfileobj = None
181
182 def __init__(self, filename=None, mode=None,
183 compresslevel=_COMPRESS_LEVEL_BEST, fileobj=None, mtime=None):
184 """Constructor for the GzipFile class.
185
186 At least one of fileobj and filename must be given a
187 non-trivial value.
188
189 The new class instance is based on fileobj, which can be a regular
190 file, an io.BytesIO object, or any other object which simulates a file.
191 It defaults to None, in which case filename is opened to provide
192 a file object.
193
194 When fileobj is not None, the filename argument is only used to be
195 included in the gzip file header, which may include the original
196 filename of the uncompressed file. It defaults to the filename of
197 fileobj, if discernible; otherwise, it defaults to the empty string,
198 and in this case the original filename is not included in the header.
199
200 The mode argument can be any of 'r', 'rb', 'a', 'ab', 'w', 'wb', 'x', or
201 'xb' depending on whether the file will be read or written. The default
202 is the mode of fileobj if discernible; otherwise, the default is 'rb'.
203 A mode of 'r' is equivalent to one of 'rb', and similarly for 'w' and
204 'wb', 'a' and 'ab', and 'x' and 'xb'.
205
206 The compresslevel argument is an integer from 0 to 9 controlling the
207 level of compression; 1 is fastest and produces the least compression,
208 and 9 is slowest and produces the most compression. 0 is no compression
209 at all. The default is 9.
210
211 The optional mtime argument is the timestamp requested by gzip. The time
212 is in Unix format, i.e., seconds since 00:00:00 UTC, January 1, 1970.
213 If mtime is omitted or None, the current time is used. Use mtime = 0
214 to generate a compressed stream that does not depend on creation time.
215
216 """
217
218 if mode and ('t' in mode or 'U' in mode):
219 raise ValueError("Invalid mode: {!r}".format(mode))
220 if mode and 'b' not in mode:
221 mode += 'b'
222
223 try:
224 if fileobj is None:
225 fileobj = self.myfileobj = builtins.open(filename, mode or 'rb')
226 if filename is None:
227 filename = getattr(fileobj, 'name', '')
228 if not isinstance(filename, (str, bytes)):
229 filename = ''
230 else:
231 filename = os.fspath(filename)
232 origmode = mode
233 if mode is None:
234 mode = getattr(fileobj, 'mode', 'rb')
235
236
237 if mode.startswith('r'):
238 self.mode = READ
239 raw = _GzipReader(fileobj)
240 self._buffer = io.BufferedReader(raw)
241 self.name = filename
242
243 elif mode.startswith(('w', 'a', 'x')):
244 if origmode is None:
245 import warnings
246 warnings.warn(
247 "GzipFile was opened for writing, but this will "
248 "change in future Python releases. "
249 "Specify the mode argument for opening it for writing.",
250 FutureWarning, 2)
251 self.mode = WRITE
252 self._init_write(filename)
253 self.compress = zlib.compressobj(compresslevel,
254 zlib.DEFLATED,
255 -zlib.MAX_WBITS,
256 zlib.DEF_MEM_LEVEL,
257 0)
258 self._write_mtime = mtime
259 self._buffer_size = _WRITE_BUFFER_SIZE
260 self._buffer = io.BufferedWriter(_WriteBufferStream(self),
261 buffer_size=self._buffer_size)
262 else:
263 raise ValueError("Invalid mode: {!r}".format(mode))
264
265 self.fileobj = fileobj
266
267 if self.mode == WRITE:
268 self._write_gzip_header(compresslevel)
269 except:
270 # Avoid a ResourceWarning if the write fails,
271 # eg read-only file or KeyboardInterrupt
272 self._close()
273 raise
274
275 @property
276 def mtime(self):
277 """Last modification time read from stream, or None"""
278 return self._buffer.raw._last_mtime
279
280 def __repr__(self):
281 s = repr(self.fileobj)
282 return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
283
284 def _init_write(self, filename):
285 self.name = filename
286 self.crc = zlib.crc32(b"")
287 self.size = 0
288 self.writebuf = []
289 self.bufsize = 0
290 self.offset = 0 # Current file offset for seek(), tell(), etc
291
292 def tell(self):
293 self._check_not_closed()
294 self._buffer.flush()
295 return super().tell()
296
297 def _write_gzip_header(self, compresslevel):
298 self.fileobj.write(b'\037\213') # magic header
299 self.fileobj.write(b'\010') # compression method
300 try:
301 # RFC 1952 requires the FNAME field to be Latin-1. Do not
302 # include filenames that cannot be represented that way.
303 fname = os.path.basename(self.name)
304 if not isinstance(fname, bytes):
305 fname = fname.encode('latin-1')
306 if fname.endswith(b'.gz'):
307 fname = fname[:-3]
308 except UnicodeEncodeError:
309 fname = b''
310 flags = 0
311 if fname:
312 flags = FNAME
313 self.fileobj.write(chr(flags).encode('latin-1'))
314 mtime = self._write_mtime
315 if mtime is None:
316 mtime = time.time()
317 write32u(self.fileobj, int(mtime))
318 if compresslevel == _COMPRESS_LEVEL_BEST:
319 xfl = b'\002'
320 elif compresslevel == _COMPRESS_LEVEL_FAST:
321 xfl = b'\004'
322 else:
323 xfl = b'\000'
324 self.fileobj.write(xfl)
325 self.fileobj.write(b'\377')
326 if fname:
327 self.fileobj.write(fname + b'\000')
328
329 def write(self,data):
330 self._check_not_closed()
331 if self.mode != WRITE:
332 import errno
333 raise OSError(errno.EBADF, "write() on read-only GzipFile object")
334
335 if self.fileobj is None:
336 raise ValueError("write() on closed GzipFile object")
337
338 return self._buffer.write(data)
339
340 def _write_raw(self, data):
341 # Called by our self._buffer underlying WriteBufferStream.
342 if isinstance(data, (bytes, bytearray)):
343 length = len(data)
344 else:
345 # accept any data that supports the buffer protocol
346 data = memoryview(data)
347 length = data.nbytes
348
349 if length > 0:
350 self.fileobj.write(self.compress.compress(data))
351 self.size += length
352 self.crc = zlib.crc32(data, self.crc)
353 self.offset += length
354
355 return length
356
357 def read(self, size=-1):
358 self._check_not_closed()
359 if self.mode != READ:
360 import errno
361 raise OSError(errno.EBADF, "read() on write-only GzipFile object")
362 return self._buffer.read(size)
363
364 def read1(self, size=-1):
365 """Implements BufferedIOBase.read1()
366
367 Reads up to a buffer's worth of data if size is negative."""
368 self._check_not_closed()
369 if self.mode != READ:
370 import errno
371 raise OSError(errno.EBADF, "read1() on write-only GzipFile object")
372
373 if size < 0:
374 size = io.DEFAULT_BUFFER_SIZE
375 return self._buffer.read1(size)
376
377 def peek(self, n):
378 self._check_not_closed()
379 if self.mode != READ:
380 import errno
381 raise OSError(errno.EBADF, "peek() on write-only GzipFile object")
382 return self._buffer.peek(n)
383
384 @property
385 def closed(self):
386 return self.fileobj is None
387
388 def close(self):
389 fileobj = self.fileobj
390 if fileobj is None or self._buffer.closed:
391 return
392 try:
393 if self.mode == WRITE:
394 self._buffer.flush()
395 fileobj.write(self.compress.flush())
396 write32u(fileobj, self.crc)
397 # self.size may exceed 2 GiB, or even 4 GiB
398 write32u(fileobj, self.size & 0xffffffff)
399 elif self.mode == READ:
400 self._buffer.close()
401 finally:
402 self._close()
403
404 def _close(self):
405 self.fileobj = None
406 myfileobj = self.myfileobj
407 if myfileobj is not None:
408 self.myfileobj = None
409 myfileobj.close()
410
411 def flush(self,zlib_mode=zlib.Z_SYNC_FLUSH):
412 self._check_not_closed()
413 if self.mode == WRITE:
414 self._buffer.flush()
415 # Ensure the compressor's buffer is flushed
416 self.fileobj.write(self.compress.flush(zlib_mode))
417 self.fileobj.flush()
418
419 def fileno(self):
420 """Invoke the underlying file object's fileno() method.
421
422 This will raise AttributeError if the underlying file object
423 doesn't support fileno().
424 """
425 return self.fileobj.fileno()
426
427 def rewind(self):
428 '''Return the uncompressed stream file position indicator to the
429 beginning of the file'''
430 if self.mode != READ:
431 raise OSError("Can't rewind in write mode")
432 self._buffer.seek(0)
433
434 def readable(self):
435 return self.mode == READ
436
437 def writable(self):
438 return self.mode == WRITE
439
440 def seekable(self):
441 return True
442
443 def seek(self, offset, whence=io.SEEK_SET):
444 if self.mode == WRITE:
445 self._check_not_closed()
446 # Flush buffer to ensure validity of self.offset
447 self._buffer.flush()
448 if whence != io.SEEK_SET:
449 if whence == io.SEEK_CUR:
450 offset = self.offset + offset
451 else:
452 raise ValueError('Seek from end not supported')
453 if offset < self.offset:
454 raise OSError('Negative seek in write mode')
455 count = offset - self.offset
456 chunk = b'\0' * self._buffer_size
457 for i in range(count // self._buffer_size):
458 self.write(chunk)
459 self.write(b'\0' * (count % self._buffer_size))
460 elif self.mode == READ:
461 self._check_not_closed()
462 return self._buffer.seek(offset, whence)
463
464 return self.offset
465
466 def readline(self, size=-1):
467 self._check_not_closed()
468 return self._buffer.readline(size)
469
470
471 def _read_exact(fp, n):
472 '''Read exactly *n* bytes from `fp`
473
474 This method is required because fp may be unbuffered,
475 i.e. return short reads.
476 '''
477 data = fp.read(n)
478 while len(data) < n:
479 b = fp.read(n - len(data))
480 if not b:
481 raise EOFError("Compressed file ended before the "
482 "end-of-stream marker was reached")
483 data += b
484 return data
485
486
487 def _read_gzip_header(fp):
488 '''Read a gzip header from `fp` and progress to the end of the header.
489
490 Returns last mtime if header was present or None otherwise.
491 '''
492 magic = fp.read(2)
493 if magic == b'':
494 return None
495
496 if magic != b'\037\213':
497 raise BadGzipFile('Not a gzipped file (%r)' % magic)
498
499 (method, flag, last_mtime) = struct.unpack("<BBIxx", _read_exact(fp, 8))
500 if method != 8:
501 raise BadGzipFile('Unknown compression method')
502
503 if flag & FEXTRA:
504 # Read & discard the extra field, if present
505 extra_len, = struct.unpack("<H", _read_exact(fp, 2))
506 _read_exact(fp, extra_len)
507 if flag & FNAME:
508 # Read and discard a null-terminated string containing the filename
509 while True:
510 s = fp.read(1)
511 if not s or s==b'\000':
512 break
513 if flag & FCOMMENT:
514 # Read and discard a null-terminated string containing a comment
515 while True:
516 s = fp.read(1)
517 if not s or s==b'\000':
518 break
519 if flag & FHCRC:
520 _read_exact(fp, 2) # Read & discard the 16-bit header CRC
521 return last_mtime
522
523
524 class _GzipReader(_compression.DecompressReader):
525 def __init__(self, fp):
526 super().__init__(_PaddedFile(fp), zlib._ZlibDecompressor,
527 wbits=-zlib.MAX_WBITS)
528 # Set flag indicating start of a new member
529 self._new_member = True
530 self._last_mtime = None
531 self._cur_block = None
532 self._next_block = 0
533
534 def _init_read(self):
535 self._crc = zlib.crc32(b"")
536 self._stream_size = 0 # Decompressed size of unconcatenated stream
537
538 def _read_gzip_header(self):
539 last_mtime = _read_gzip_header(self._fp)
540 if last_mtime is None:
541 return False
542 self._last_mtime = last_mtime
543 return True
544
545 def read(self, size=-1):
546 if size < 0:
547 return self.readall()
548 # size=0 is special because decompress(max_length=0) is not supported
549 if not size:
550 return b""
551
552 # For certain input data, a single
553 # call to decompress() may not return
554 # any data. In this case, retry until we get some data or reach EOF.
555 while True:
556 if self._decompressor.eof:
557 # Ending case: we've come to the end of a member in the file,
558 # so finish up this member, and read a new gzip header.
559 # Check the CRC and file size, and set the flag so we read
560 # a new member
561 self._read_eof()
562 self._new_member = True
563 self._decompressor = self._decomp_factory(
564 **self._decomp_args)
565
566 if self._new_member:
567 # If the _new_member flag is set, we have to
568 # jump to the next member, if there is one.
569 self._init_read()
570 if not self._read_gzip_header():
571 return b""
572 self._new_member = False
573
574 # Read a chunk of data from the file
575 if self._decompressor.needs_input:
576 buf = self._fp.read(READ_BUFFER_SIZE)
577 uncompress = self._decompressor.decompress(buf, size)
578 else:
579 uncompress = self._decompressor.decompress(b"", size)
580
581 if self._decompressor.unused_data != b"":
582 # Prepend the already read bytes to the fileobj so they can
583 # be seen by _read_eof() and _read_gzip_header()
584 self._fp.prepend(self._decompressor.unused_data)
585
586 if uncompress != b"":
587 break
588 if buf == b"":
589 raise EOFError("Compressed file ended before the "
590 "end-of-stream marker was reached")
591
592 self._crc = zlib.crc32(uncompress, self._crc)
593 self._stream_size += len(uncompress)
594 return uncompress
595
596 def _read_eof(self):
597 # We've read to the end of the file
598 # We check that the computed CRC and size of the
599 # uncompressed data matches the stored values. Note that the size
600 # stored is the true file size mod 2**32.
601 crc32, isize = struct.unpack("<II", _read_exact(self._fp, 8))
602 if crc32 != self._crc:
603 raise BadGzipFile("CRC check failed %s != %s" % (hex(crc32),
604 hex(self._crc)))
605 elif isize != (self._stream_size & 0xffffffff):
606 raise BadGzipFile("Incorrect length of data produced")
607
608 # Gzip files can be padded with zeroes and still have archives.
609 # Consume all zero bytes and set the file position to the first
610 # non-zero byte. See http://www.gzip.org/#faq8
611 c = b"\x00"
612 while c == b"\x00":
613 c = self._fp.read(1)
614 # At this point we know how big the compressed data was:
615 # 2+8+[x-y]+8+(1 * nzeros), which is available via _fp.tell()
616 if c:
617 self._fp.prepend(c)
618
619 self._cur_block = self._next_block
620 self._next_block = self._fp.tell()
621
622 def _rewind(self):
623 super()._rewind()
624 self._new_member = True
625
626
627 def compress(data, compresslevel=_COMPRESS_LEVEL_BEST, *, mtime=None):
628 """Compress data in one shot and return the compressed string.
629
630 compresslevel sets the compression level in range of 0-9.
631 mtime can be used to set the modification time. The modification time is
632 set to the current time by default.
633 """
634 # Wbits=31 automatically includes a gzip header and trailer.
635 gzip_data = zlib.compress(data, level=compresslevel, wbits=31)
636 if mtime is None:
637 mtime = time.time()
638 # Reuse gzip header created by zlib, replace mtime and OS byte for
639 # consistency.
640 header = struct.pack("<4sLBB", gzip_data, int(mtime), gzip_data[8], 255)
641 return header + gzip_data[10:]
642
643
644 def decompress(data):
645 """Decompress a gzip compressed string in one shot.
646 Return the decompressed string.
647 """
648 decompressed_members = []
649 while True:
650 fp = io.BytesIO(data)
651 if _read_gzip_header(fp) is None:
652 return b"".join(decompressed_members)
653 # Use a zlib raw deflate compressor
654 do = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
655 # Read all the data except the header
656 decompressed = do.decompress(data[fp.tell():])
657 if not do.eof or len(do.unused_data) < 8:
658 raise EOFError("Compressed file ended before the end-of-stream "
659 "marker was reached")
660 crc, length = struct.unpack("<II", do.unused_data[:8])
661 if crc != zlib.crc32(decompressed):
662 raise BadGzipFile("CRC check failed")
663 if length != (len(decompressed) & 0xffffffff):
664 raise BadGzipFile("Incorrect length of data produced")
665 decompressed_members.append(decompressed)
666 data = do.unused_data[8:].lstrip(b"\x00")
667
668 def main():
669 from argparse import ArgumentParser
670 parser = ArgumentParser(description=
671 "A simple command line interface for the gzip module: act like gzip, "
672 "but do not delete the input file.")
673 group = parser.add_mutually_exclusive_group()
674 group.add_argument('--fast', action='store_true', help='compress faster')
675 group.add_argument('--best', action='store_true', help='compress better')
676 group.add_argument("-d", "--decompress", action="store_true",
677 help="act like gunzip instead of gzip")
678
679 parser.add_argument("args", nargs="*", default=["-"], metavar='file')
680 args = parser.parse_args()
681
682 compresslevel = _COMPRESS_LEVEL_TRADEOFF
683 if args.fast:
684 compresslevel = _COMPRESS_LEVEL_FAST
685 elif args.best:
686 compresslevel = _COMPRESS_LEVEL_BEST
687
688 for arg in args.args:
689 if args.decompress:
690 if arg == "-":
691 f = GzipFile(filename="", mode="rb", fileobj=sys.stdin.buffer)
692 g = sys.stdout.buffer
693 else:
694 if arg[-3:] != ".gz":
695 sys.exit(f"filename doesn't end in .gz: {arg!r}")
696 f = open(arg, "rb")
697 g = builtins.open(arg[:-3], "wb")
698 else:
699 if arg == "-":
700 f = sys.stdin.buffer
701 g = GzipFile(filename="", mode="wb", fileobj=sys.stdout.buffer,
702 compresslevel=compresslevel)
703 else:
704 f = builtins.open(arg, "rb")
705 g = open(arg + ".gz", "wb")
706 while True:
707 chunk = f.read(READ_BUFFER_SIZE)
708 if not chunk:
709 break
710 g.write(chunk)
711 if g is not sys.stdout.buffer:
712 g.close()
713 if f is not sys.stdin.buffer:
714 f.close()
715
716 if __name__ == '__main__':
717 main()