Lib/gzip.py

import time
import string
import zlib
import StringIO
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(buf) + self.extrasize

    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])
		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()