0
0
mirror of https://github.com/python/cpython.git synced 2024-12-01 03:01:36 +01:00
cpython/Lib/gzip.py
Jeremy Hylton ee918cb487 Fix bug reported by Harri Pasanen: gzip + cPickle doesn't work. The
problem was a couple of bugs in the readline implementation.

1. Include the '\n' in the string returned by readline
2. Bug calculating new buffer size in _unread

Also remove unncessary import of StringIO
1998-05-13 21:49:58 +00:00

308 lines
9.3 KiB
Python

import time
import string
import zlib
import __builtin__
# implements a python function that reads and writes a gzipped file
# the user of the file doesn't have to worry about the compression,
# but random access is not allowed
# based on Andrew Kuchling's minigzip.py distributed with the zlib module
FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
READ, WRITE = 1, 2
def write32(output, value):
t = divmod(value, 256)
b1 = chr(t[1])
t = divmod(t[0], 256)
b2 = chr(t[1])
t = divmod(t[0], 256)
b3 = chr(t[1])
t = divmod(t[0], 256)
b4 = chr(t[1])
buf = b1 + b2 + b3 + b4
output.write(buf)
def read32(input):
buf = input.read(4)
v = ord(buf[0])
v = v + (ord(buf[1]) << 8)
v = v + (ord(buf[2]) << 16)
v = v + (ord(buf[3]) << 24)
return v
def open(filename, mode="r", compresslevel=9):
return GzipFile(filename, mode, compresslevel)
class GzipFile:
myfileobj = None
def __init__(self, filename=None, mode=None,
compresslevel=9, fileobj=None):
if fileobj is None:
fileobj = self.myfileobj = __builtin__.open(filename, mode or 'r')
if filename is None:
if hasattr(fileobj, 'name'): filename = fileobj.name
else: filename = ''
if mode is None:
if hasattr(fileobj, 'mode'): mode = fileobj.mode
else: mode = 'r'
if mode[0:1] == 'r':
self.mode = READ
self._init_read()
self.filename = filename
self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
elif mode[0:1] == 'w':
self.mode = WRITE
self._init_write(filename)
self.compress = zlib.compressobj(compresslevel,
zlib.DEFLATED,
-zlib.MAX_WBITS,
zlib.DEF_MEM_LEVEL,
0)
else:
raise ValueError, "Mode " + mode + " not supported"
self.fileobj = fileobj
if self.mode == WRITE:
self._write_gzip_header()
elif self.mode == READ:
self._read_gzip_header()
def __repr__(self):
s = repr(self.fileobj)
return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
def _init_write(self, filename):
if filename[-3:] != '.gz':
filename = filename + '.gz'
self.filename = filename
self.crc = zlib.crc32("")
self.size = 0
self.writebuf = []
self.bufsize = 0
def _write_gzip_header(self):
self.fileobj.write('\037\213') # magic header
self.fileobj.write('\010') # compression method
fname = self.filename[:-3]
flags = 0
if fname:
flags = FNAME
self.fileobj.write(chr(flags))
write32(self.fileobj, int(time.time()))
self.fileobj.write('\002')
self.fileobj.write('\377')
if fname:
self.fileobj.write(fname + '\000')
def _init_read(self):
self.crc = zlib.crc32("")
self.size = 0
self.extrabuf = ""
self.extrasize = 0
def _read_gzip_header(self):
magic = self.fileobj.read(2)
if magic != '\037\213':
raise RuntimeError, 'Not a gzipped file'
method = ord( self.fileobj.read(1) )
if method != 8:
raise RuntimeError, 'Unknown compression method'
flag = ord( self.fileobj.read(1) )
# modtime = self.fileobj.read(4)
# extraflag = self.fileobj.read(1)
# os = self.fileobj.read(1)
self.fileobj.read(6)
if flag & FEXTRA:
# Read & discard the extra field, if present
xlen=ord(self.fileobj.read(1))
xlen=xlen+256*ord(self.fileobj.read(1))
self.fileobj.read(xlen)
if flag & FNAME:
# Read and discard a null-terminated string containing the filename
while (1):
s=self.fileobj.read(1)
if not s or s=='\000': break
if flag & FCOMMENT:
# Read and discard a null-terminated string containing a comment
while (1):
s=self.fileobj.read(1)
if not s or s=='\000': break
if flag & FHCRC:
self.fileobj.read(2) # Read & discard the 16-bit header CRC
def write(self,data):
if self.fileobj is None:
raise ValueError, "write() on closed GzipFile object"
if len(data) > 0:
self.size = self.size + len(data)
self.crc = zlib.crc32(data, self.crc)
self.fileobj.write( self.compress.compress(data) )
def writelines(self,lines):
self.write(string.join(lines))
def read(self, size=None):
if self.extrasize <= 0 and self.fileobj is None:
return ''
readsize = 1024
if not size: # get the whole thing
try:
while 1:
self._read(readsize)
readsize = readsize * 2
except EOFError:
size = self.extrasize
else: # just get some more of it
try:
while size > self.extrasize:
self._read(readsize)
readsize = readsize * 2
except EOFError:
pass
chunk = self.extrabuf[:size]
self.extrabuf = self.extrabuf[size:]
self.extrasize = self.extrasize - size
return chunk
def _unread(self, buf):
self.extrabuf = buf + self.extrabuf
self.extrasize = len(self.extrabuf)
def _read(self, size=1024):
try:
buf = self.fileobj.read(size)
except AttributeError:
raise EOFError, "Reached EOF"
if buf == "":
uncompress = self.decompress.flush()
if uncompress == "":
self._read_eof()
self.fileobj = None
raise EOFError, 'Reached EOF'
else:
uncompress = self.decompress.decompress(buf)
self.crc = zlib.crc32(uncompress, self.crc)
self.extrabuf = self.extrabuf + uncompress
self.extrasize = self.extrasize + len(uncompress)
self.size = self.size + len(uncompress)
def _read_eof(self):
# Andrew writes:
## We've read to the end of the file, so we have to rewind in order
## to reread the 8 bytes containing the CRC and the file size. The
## decompressor is smart and knows when to stop, so feeding it
## extra data is harmless.
self.fileobj.seek(-8, 2)
crc32 = read32(self.fileobj)
isize = read32(self.fileobj)
if crc32 != self.crc:
self.error = "CRC check failed"
elif isize != self.size:
self.error = "Incorrect length of data produced"
def close(self):
if self.mode == WRITE:
self.fileobj.write(self.compress.flush())
write32(self.fileobj, self.crc)
write32(self.fileobj, self.size)
self.fileobj = None
elif self.mode == READ:
self.fileobj = None
if self.myfileobj:
self.myfileobj.close()
self.myfileobj = None
def flush(self):
self.fileobj.flush()
def seek(self):
raise IOError, 'Random access not allowed in gzip files'
def tell(self):
raise IOError, 'I won\'t tell() you for gzip files'
def isatty(self):
return 0
def readline(self):
bufs = []
readsize = 100
while 1:
c = self.read(readsize)
i = string.find(c, '\n')
if i >= 0 or c == '':
bufs.append(c[:i+1])
self._unread(c[i+1:])
return string.join(bufs, '')
bufs.append(c)
readsize = readsize * 2
def readlines(self):
buf = self.read()
return string.split(buf, '\n')
def writelines(self, L):
for line in L:
self.write(line)
def _test():
# Act like gzip; with -d, act like gunzip.
# The input file is not deleted, however, nor are any other gzip
# options or features supported.
import sys
args = sys.argv[1:]
decompress = args and args[0] == "-d"
if decompress:
args = args[1:]
if not args:
args = ["-"]
for arg in args:
if decompress:
if arg == "-":
f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
g = sys.stdout
else:
if arg[-3:] != ".gz":
print "filename doesn't end in .gz:", `arg`
continue
f = open(arg, "rb")
g = __builtin__.open(arg[:-3], "wb")
else:
if arg == "-":
f = sys.stdin
g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
else:
f = __builtin__.open(arg, "rb")
g = open(arg + ".gz", "wb")
while 1:
chunk = f.read(1024)
if not chunk:
break
g.write(chunk)
if g is not sys.stdout:
g.close()
if f is not sys.stdin:
f.close()
if __name__ == '__main__':
_test()