Lib/gzip.py

   1 import time
   2 import string
   3 import zlib
   4 import struct
   5 import __builtin__
   6
   7 # implements a python function that reads and writes a gzipped file
   8 # the user of the file doesn't have to worry about the compression,
   9 # but random access is not allowed
  10
  11 # based on Andrew Kuchling's minigzip.py distributed with the zlib module
  12
  13 FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
  14
  15 READ, WRITE = 1, 2
  16
  17 def write32(output, value):
  18     output.write(struct.pack("<l", value))
  19
  20 def write32u(output, value):
  21     output.write(struct.pack("<L", value))
  22
  23 def read32(input):
  24     return struct.unpack("<l", input.read(4))[0]
  25
  26 def open(filename, mode="rb", compresslevel=9):
  27     return GzipFile(filename, mode, compresslevel)
  28
  29 class GzipFile:
  30
  31     myfileobj = None
  32
  33     def __init__(self, filename=None, mode=None,
  34                  compresslevel=9, fileobj=None):
  35         if fileobj is None:
  36             fileobj = self.myfileobj = __builtin__.open(filename, mode or 'rb')
  37         if filename is None:
  38             if hasattr(fileobj, 'name'): filename = fileobj.name
  39             else: filename = ''
  40         if mode is None:
  41             if hasattr(fileobj, 'mode'): mode = fileobj.mode
  42             else: mode = 'rb'
  43
  44         if mode[0:1] == 'r':
  45             self.mode = READ
  46             # Set flag indicating start of a new member
  47             self._new_member = 1
  48             self.extrabuf = ""
  49             self.extrasize = 0
  50             self.filename = filename
  51
  52         elif mode[0:1] == 'w' or mode[0:1] == 'a':
  53             self.mode = WRITE
  54             self._init_write(filename)
  55             self.compress = zlib.compressobj(compresslevel,
  56                                              zlib.DEFLATED,
  57                                              -zlib.MAX_WBITS,
  58                                              zlib.DEF_MEM_LEVEL,
  59                                              0)
  60         else:
  61             raise ValueError, "Mode " + mode + " not supported"
  62
  63         self.fileobj = fileobj
  64
  65         if self.mode == WRITE:
  66             self._write_gzip_header()
  67
  68     def __repr__(self):
  69         s = repr(self.fileobj)
  70         return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
  71
  72     def _init_write(self, filename):
  73         if filename[-3:] != '.gz':
  74             filename = filename + '.gz'
  75         self.filename = filename
  76         self.crc = zlib.crc32("")
  77         self.size = 0
  78         self.writebuf = []
  79         self.bufsize = 0
  80
  81     def _write_gzip_header(self):
  82         self.fileobj.write('\037\213')             # magic header
  83         self.fileobj.write('\010')                 # compression method
  84         fname = self.filename[:-3]
  85         flags = 0
  86         if fname:
  87             flags = FNAME
  88         self.fileobj.write(chr(flags))
  89         write32u(self.fileobj, long(time.time()))
  90         self.fileobj.write('\002')
  91         self.fileobj.write('\377')
  92         if fname:
  93             self.fileobj.write(fname + '\000')
  94
  95     def _init_read(self):
  96         self.crc = zlib.crc32("")
  97         self.size = 0
  98
  99     def _read_gzip_header(self):
 100         magic = self.fileobj.read(2)
 101         if magic != '\037\213':
 102             raise IOError, 'Not a gzipped file'
 103         method = ord( self.fileobj.read(1) )
 104         if method != 8:
 105             raise IOError, 'Unknown compression method'
 106         flag = ord( self.fileobj.read(1) )
 107         # modtime = self.fileobj.read(4)
 108         # extraflag = self.fileobj.read(1)
 109         # os = self.fileobj.read(1)
 110         self.fileobj.read(6)
 111
 112         if flag & FEXTRA:
 113             # Read & discard the extra field, if present
 114             xlen=ord(self.fileobj.read(1))
 115             xlen=xlen+256*ord(self.fileobj.read(1))
 116             self.fileobj.read(xlen)
 117         if flag & FNAME:
 118             # Read and discard a null-terminated string containing the filename
 119             while (1):
 120                 s=self.fileobj.read(1)
 121                 if not s or s=='\000': break
 122         if flag & FCOMMENT:
 123             # Read and discard a null-terminated string containing a comment
 124             while (1):
 125                 s=self.fileobj.read(1)
 126                 if not s or s=='\000': break
 127         if flag & FHCRC:
 128             self.fileobj.read(2)     # Read & discard the 16-bit header CRC
 129
 130
 131     def write(self,data):
 132         if self.fileobj is None:
 133             raise ValueError, "write() on closed GzipFile object"
 134         if len(data) > 0:
 135             self.size = self.size + len(data)
 136             self.crc = zlib.crc32(data, self.crc)
 137             self.fileobj.write( self.compress.compress(data) )
 138
 139     def writelines(self,lines):
 140         self.write(string.join(lines))
 141
 142     def read(self, size=None):
 143         if self.extrasize <= 0 and self.fileobj is None:
 144             return ''
 145
 146         readsize = 1024
 147         if not size:        # get the whole thing
 148             try:
 149                 while 1:
 150                     self._read(readsize)
 151                     readsize = readsize * 2
 152             except EOFError:
 153                 size = self.extrasize
 154         else:               # just get some more of it
 155             try:
 156                 while size > self.extrasize:
 157                     self._read(readsize)
 158                     readsize = readsize * 2
 159             except EOFError:
 160                 if size > self.extrasize:
 161                     size = self.extrasize
 162
 163         chunk = self.extrabuf[:size]
 164         self.extrabuf = self.extrabuf[size:]
 165         self.extrasize = self.extrasize - size
 166
 167         return chunk
 168
 169     def _unread(self, buf):
 170         self.extrabuf = buf + self.extrabuf
 171         self.extrasize = len(buf) + self.extrasize
 172
 173     def _read(self, size=1024):
 174         if self.fileobj is None: raise EOFError, "Reached EOF"
 175
 176         if self._new_member:
 177             # If the _new_member flag is set, we have to
 178             #
 179             # First, check if we're at the end of the file;
 180             # if so, it's time to stop; no more members to read.
 181             pos = self.fileobj.tell()   # Save current position
 182             self.fileobj.seek(0, 2)     # Seek to end of file
 183             if pos == self.fileobj.tell():
 184                 self.fileobj = None
 185                 raise EOFError, "Reached EOF"
 186             else:
 187                 self.fileobj.seek( pos ) # Return to original position
 188
 189             self._init_read()
 190             self._read_gzip_header()
 191             self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
 192             self._new_member = 0
 193
 194         # Read a chunk of data from the file
 195         buf = self.fileobj.read(size)
 196
 197         # If the EOF has been reached, flush the decompression object
 198         # and mark this object as finished.
 199
 200         if buf == "":
 201             uncompress = self.decompress.flush()
 202             self._read_eof()
 203             self.fileobj = None
 204             self._add_read_data( uncompress )
 205             raise EOFError, 'Reached EOF'
 206
 207         uncompress = self.decompress.decompress(buf)
 208         self._add_read_data( uncompress )
 209
 210         if self.decompress.unused_data != "":
 211             # Ending case: we've come to the end of a member in the file,
 212             # so seek back to the start of the unused data, finish up
 213             # this member, and read a new gzip header.
 214             # (The number of bytes to seek back is the length of the unused
 215             # data, minus 8 because _read_eof() will rewind a further 8 bytes)
 216             self.fileobj.seek( -len(self.decompress.unused_data)+8, 1)
 217
 218             # Check the CRC and file size, and set the flag so we read
 219             # a new member on the next call
 220             self._read_eof()
 221             self._new_member = 1
 222
 223     def _add_read_data(self, data):
 224         self.crc = zlib.crc32(data, self.crc)
 225         self.extrabuf = self.extrabuf + data
 226         self.extrasize = self.extrasize + len(data)
 227         self.size = self.size + len(data)
 228
 229     def _read_eof(self):
 230         # We've read to the end of the file, so we have to rewind in order
 231         # to reread the 8 bytes containing the CRC and the file size.
 232         # We check the that the computed CRC and size of the
 233         # uncompressed data matches the stored values.
 234         self.fileobj.seek(-8, 1)
 235         crc32 = read32(self.fileobj)
 236         isize = read32(self.fileobj)
 237         if crc32%0x100000000L != self.crc%0x100000000L:
 238             raise ValueError, "CRC check failed"
 239         elif isize != self.size:
 240             raise ValueError, "Incorrect length of data produced"
 241
 242     def close(self):
 243         if self.mode == WRITE:
 244             self.fileobj.write(self.compress.flush())
 245             write32(self.fileobj, self.crc)
 246             write32(self.fileobj, self.size)
 247             self.fileobj = None
 248         elif self.mode == READ:
 249             self.fileobj = None
 250         if self.myfileobj:
 251             self.myfileobj.close()
 252             self.myfileobj = None
 253
 254     def __del__(self):
 255         if (self.myfileobj is not None or
 256             self.fileobj is not None):
 257             self.close()
 258
 259     def flush(self):
 260         self.fileobj.flush()
 261
 262     def seek(self):
 263         raise IOError, 'Random access not allowed in gzip files'
 264
 265     def tell(self):
 266         raise IOError, 'I won\'t tell() you for gzip files'
 267
 268     def isatty(self):
 269         return 0
 270
 271     def readline(self):
 272         bufs = []
 273         readsize = 100
 274         while 1:
 275             c = self.read(readsize)
 276             i = string.find(c, '\n')
 277             if i >= 0 or c == '':
 278                 bufs.append(c[:i+1])
 279                 self._unread(c[i+1:])
 280                 return string.join(bufs, '')
 281             bufs.append(c)
 282             readsize = readsize * 2
 283
 284     def readlines(self):
 285         buf = self.read()
 286         lines = string.split(buf, '\n')
 287         for i in range(len(lines)-1):
 288             lines[i] = lines[i] + '\n'
 289         if lines and not lines[-1]:
 290             del lines[-1]
 291         return lines
 292
 293     def writelines(self, L):
 294         for line in L:
 295             self.write(line)
 296
 297
 298 def _test():
 299     # Act like gzip; with -d, act like gunzip.
 300     # The input file is not deleted, however, nor are any other gzip
 301     # options or features supported.
 302     import sys
 303     args = sys.argv[1:]
 304     decompress = args and args[0] == "-d"
 305     if decompress:
 306         args = args[1:]
 307     if not args:
 308         args = ["-"]
 309     for arg in args:
 310         if decompress:
 311             if arg == "-":
 312                 f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
 313                 g = sys.stdout
 314             else:
 315                 if arg[-3:] != ".gz":
 316                     print "filename doesn't end in .gz:", `arg`
 317                     continue
 318                 f = open(arg, "rb")
 319                 g = __builtin__.open(arg[:-3], "wb")
 320         else:
 321             if arg == "-":
 322                 f = sys.stdin
 323                 g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
 324             else:
 325                 f = __builtin__.open(arg, "rb")
 326                 g = open(arg + ".gz", "wb")
 327         while 1:
 328             chunk = f.read(1024)
 329             if not chunk:
 330                 break
 331             g.write(chunk)
 332         if g is not sys.stdout:
 333             g.close()
 334         if f is not sys.stdin:
 335             f.close()
 336
 337 if __name__ == '__main__':
 338     _test()