Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 1 | import time |
| 2 | import string |
| 3 | import zlib |
| 4 | import StringIO |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 5 | import __builtin__ |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 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 sequential 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 | t = divmod(value, 256) |
| 19 | b1 = chr(t[1]) |
| 20 | |
| 21 | t = divmod(t[0], 256) |
| 22 | b2 = chr(t[1]) |
| 23 | |
| 24 | t = divmod(t[0], 256) |
| 25 | b3 = chr(t[1]) |
| 26 | |
| 27 | t = divmod(t[0], 256) |
| 28 | b4 = chr(t[1]) |
| 29 | |
| 30 | buf = b1 + b2 + b3 + b4 |
| 31 | output.write(buf) |
| 32 | |
| 33 | |
| 34 | def read32(input): |
| 35 | buf = input.read(4) |
| 36 | v = ord(buf[0]) |
| 37 | v = v + (ord(buf[1]) << 8) |
| 38 | v = v + (ord(buf[2]) << 16) |
| 39 | v = v + (ord(buf[3]) << 24) |
| 40 | return v |
| 41 | |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 42 | def open(filename, mode="r", compresslevel=9): |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 43 | return GzipFile(filename, mode, compresslevel) |
| 44 | |
| 45 | class GzipFile: |
| 46 | |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 47 | myfileobj = None |
| 48 | |
| 49 | def __init__(self, filename=None, mode=None, |
| 50 | compresslevel=9, fileobj=None): |
| 51 | if fileobj is None: |
| 52 | fileobj = self.myfileobj = __builtin__.open(filename, mode or 'r') |
| 53 | if filename is None: |
| 54 | if hasattr(fileobj, 'name'): filename = fileobj.name |
| 55 | else: filename = 'GzippedFile' |
| 56 | if mode is None: |
| 57 | if hasattr(fileobj, 'mode'): mode = fileobj.mode |
| 58 | else: mode = 'r' |
| 59 | |
| 60 | if mode[0:1] == 'r': |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 61 | self.mode = READ |
| 62 | self._init_read() |
| 63 | self.filename = filename |
| 64 | self.decompress = zlib.decompressobj(-zlib.MAX_WBITS) |
| 65 | |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 66 | elif mode[0:1] == 'w': |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 67 | self.mode = WRITE |
| 68 | self._init_write(filename) |
| 69 | self.compress = zlib.compressobj(compresslevel, |
| 70 | zlib.DEFLATED, |
| 71 | -zlib.MAX_WBITS, |
| 72 | zlib.DEF_MEM_LEVEL, |
| 73 | 0) |
| 74 | else: |
| 75 | raise ValueError, "Mode " + mode + " not supported" |
| 76 | |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 77 | self.fileobj = fileobj |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 78 | |
| 79 | if self.mode == WRITE: |
| 80 | self._write_gzip_header() |
| 81 | elif self.mode == READ: |
| 82 | self._read_gzip_header() |
| 83 | |
| 84 | |
| 85 | def __repr__(self): |
| 86 | s = repr(self.fileobj) |
| 87 | return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>' |
| 88 | |
| 89 | def _init_write(self, filename): |
| 90 | if filename[-3:] != '.gz': |
| 91 | filename = filename + '.gz' |
| 92 | self.filename = filename |
| 93 | self.crc = zlib.crc32("") |
| 94 | self.size = 0 |
| 95 | self.writebuf = [] |
| 96 | self.bufsize = 0 |
| 97 | |
| 98 | def _write_gzip_header(self): |
| 99 | self.fileobj.write('\037\213') # magic header |
| 100 | self.fileobj.write('\010') # compression method |
| 101 | self.fileobj.write(chr(FNAME)) |
| 102 | write32(self.fileobj, int(time.time())) |
| 103 | self.fileobj.write('\002') |
| 104 | self.fileobj.write('\377') |
| 105 | self.fileobj.write(self.filename[:-3] + '\000') |
| 106 | |
| 107 | def _init_read(self): |
| 108 | self.crc = zlib.crc32("") |
| 109 | self.size = 0 |
| 110 | self.extrabuf = "" |
| 111 | self.extrasize = 0 |
| 112 | |
| 113 | def _read_gzip_header(self): |
| 114 | magic = self.fileobj.read(2) |
| 115 | if magic != '\037\213': |
| 116 | raise RuntimeError, 'Not a gzipped file' |
| 117 | method = ord( self.fileobj.read(1) ) |
| 118 | if method != 8: |
| 119 | raise RuntimeError, 'Unknown compression method' |
| 120 | flag = ord( self.fileobj.read(1) ) |
| 121 | # modtime = self.fileobj.read(4) |
| 122 | # extraflag = self.fileobj.read(1) |
| 123 | # os = self.fileobj.read(1) |
| 124 | self.fileobj.read(6) |
| 125 | |
| 126 | if flag & FEXTRA: |
| 127 | # Read & discard the extra field, if present |
| 128 | xlen=ord(self.fileobj.read(1)) |
| 129 | xlen=xlen+256*ord(self.fileobj.read(1)) |
| 130 | self.fileobj.read(xlen) |
| 131 | if flag & FNAME: |
| 132 | # Read and discard a null-terminated string containing the filename |
| 133 | while (1): |
| 134 | s=self.fileobj.read(1) |
| 135 | if s=='\000': break |
| 136 | if flag & FCOMMENT: |
| 137 | # Read and discard a null-terminated string containing a comment |
| 138 | while (1): |
| 139 | s=self.fileobj.read(1) |
| 140 | if s=='\000': break |
| 141 | if flag & FHCRC: |
| 142 | self.fileobj.read(2) # Read & discard the 16-bit header CRC |
| 143 | |
| 144 | |
| 145 | def write(self,data): |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 146 | if self.fileobj is None: |
| 147 | raise ValueError, "write() on closed GzipFile object" |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 148 | if len(data) > 0: |
| 149 | self.size = self.size + len(data) |
| 150 | self.crc = zlib.crc32(data, self.crc) |
| 151 | self.fileobj.write( self.compress.compress(data) ) |
| 152 | |
| 153 | def writelines(self,lines): |
| 154 | self.write(string.join(lines)) |
| 155 | |
| 156 | def read(self,size=None): |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 157 | if self.extrasize <= 0 and self.fileobj is None: |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 158 | return '' |
| 159 | |
| 160 | if not size: |
| 161 | # get the whole thing |
| 162 | try: |
| 163 | while 1: |
| 164 | self._read() |
| 165 | except EOFError: |
| 166 | size = self.extrasize |
| 167 | else: |
| 168 | # just get some more of it |
| 169 | try: |
| 170 | while size > self.extrasize: |
| 171 | self._read() |
| 172 | except EOFError: |
| 173 | pass |
| 174 | |
| 175 | chunk = self.extrabuf[:size] |
| 176 | self.extrabuf = self.extrabuf[size:] |
| 177 | self.extrasize = self.extrasize - size |
| 178 | |
| 179 | return chunk |
| 180 | |
| 181 | def _read(self): |
| 182 | buf = self.fileobj.read(1024) |
| 183 | if buf == "": |
| 184 | uncompress = self.decompress.flush() |
| 185 | if uncompress == "": |
| 186 | self._read_eof() |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 187 | self.fileobj = None |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 188 | raise EOFError, 'Reached EOF' |
| 189 | else: |
| 190 | uncompress = self.decompress.decompress(buf) |
| 191 | self.crc = zlib.crc32(uncompress, self.crc) |
| 192 | self.extrabuf = self.extrabuf + uncompress |
| 193 | self.extrasize = self.extrasize + len(uncompress) |
| 194 | self.size = self.size + len(uncompress) |
| 195 | |
| 196 | def _read_eof(self): |
| 197 | # Andrew writes: |
| 198 | ## We've read to the end of the file, so we have to rewind in order |
| 199 | ## to reread the 8 bytes containing the CRC and the file size. The |
| 200 | ## decompressor is smart and knows when to stop, so feeding it |
| 201 | ## extra data is harmless. |
| 202 | self.fileobj.seek(-8, 2) |
| 203 | crc32 = read32(self.fileobj) |
| 204 | isize = read32(self.fileobj) |
| 205 | if crc32 != self.crc: |
| 206 | self.error = "CRC check failed" |
| 207 | elif isize != self.size: |
| 208 | self.error = "Incorrect length of data produced" |
| 209 | |
| 210 | def close(self): |
| 211 | if self.mode == WRITE: |
| 212 | self.fileobj.write(self.compress.flush()) |
| 213 | write32(self.fileobj, self.crc) |
| 214 | write32(self.fileobj, self.size) |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 215 | self.fileobj = None |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 216 | elif self.mode == READ: |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 217 | self.fileobj = None |
| 218 | if self.myfileobj: |
| 219 | self.myfileobj.close() |
| 220 | self.myfileobj = None |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 221 | |
| 222 | def flush(self): |
| 223 | self.fileobj.flush() |
| 224 | |
| 225 | def seek(self): |
| 226 | raise IOError, 'Random access not allowed in gzip files' |
| 227 | |
| 228 | def tell(self): |
| 229 | raise IOError, 'I won\'t tell() you for gzip files' |
| 230 | |
| 231 | def isatty(self): |
| 232 | return 0 |
| 233 | |
| 234 | def readline(self): |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 235 | # XXX This function isn't implemented in a very efficient way |
| 236 | line="" |
| 237 | while 1: |
| 238 | c = self.read(1) |
| 239 | line = line + c |
| 240 | if c=='\n' or c=="": break |
| 241 | return line |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 242 | |
| 243 | def readlines(self): |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 244 | L=[] |
| 245 | line = self.readline() |
| 246 | while line!="": |
| 247 | L.append(line) |
| 248 | line = self.readline() |
| 249 | return L |
Guido van Rossum | 1526219 | 1997-04-30 16:04:57 +0000 | [diff] [blame] | 250 | |
Guido van Rossum | 68de379 | 1997-07-19 20:22:23 +0000 | [diff] [blame] | 251 | def writelines(self, L): |
| 252 | for line in L: |
| 253 | self.write(line) |