blob: 79ca2050d2e77ee68c761081491bf791a5d6aefb [file] [log] [blame]
Guido van Rossum4b8c6ea2000-02-04 15:39:30 +00001"""Functions that read and write gzipped files.
2
Guido van Rossum54f22ed2000-02-04 15:10:34 +00003The user of the file doesn't have to worry about the compression,
4but random access is not allowed."""
5
6# based on Andrew Kuchling's minigzip.py distributed with the zlib module
7
Eric S. Raymondee5e61d2001-02-09 09:10:35 +00008import struct, sys, time
Guido van Rossum15262191997-04-30 16:04:57 +00009import zlib
Guido van Rossum68de3791997-07-19 20:22:23 +000010import __builtin__
Guido van Rossum15262191997-04-30 16:04:57 +000011
Skip Montanaro2dd42762001-01-23 15:35:05 +000012__all__ = ["GzipFile","open"]
13
Guido van Rossum15262191997-04-30 16:04:57 +000014FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
15
16READ, WRITE = 1, 2
17
18def write32(output, value):
Jeremy Hyltonc19f9971999-03-23 23:05:34 +000019 output.write(struct.pack("<l", value))
Tim Peters07e99cb2001-01-14 23:47:14 +000020
Guido van Rossum95bdd0b1999-04-12 14:34:16 +000021def write32u(output, value):
Jack Jansen2d0589b2001-02-21 10:39:35 +000022 if value < 0:
23 value = value + 0x100000000L
Guido van Rossum95bdd0b1999-04-12 14:34:16 +000024 output.write(struct.pack("<L", value))
25
Guido van Rossum15262191997-04-30 16:04:57 +000026def read32(input):
Jeremy Hyltonc19f9971999-03-23 23:05:34 +000027 return struct.unpack("<l", input.read(4))[0]
Guido van Rossum15262191997-04-30 16:04:57 +000028
Fred Drakefa1591c1999-04-05 18:37:59 +000029def open(filename, mode="rb", compresslevel=9):
Guido van Rossum15262191997-04-30 16:04:57 +000030 return GzipFile(filename, mode, compresslevel)
31
32class GzipFile:
33
Guido van Rossum68de3791997-07-19 20:22:23 +000034 myfileobj = None
35
Tim Peters07e99cb2001-01-14 23:47:14 +000036 def __init__(self, filename=None, mode=None,
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000037 compresslevel=9, fileobj=None):
38 if fileobj is None:
Fred Drake9bb76d11999-04-05 18:33:40 +000039 fileobj = self.myfileobj = __builtin__.open(filename, mode or 'rb')
Guido van Rossum68de3791997-07-19 20:22:23 +000040 if filename is None:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000041 if hasattr(fileobj, 'name'): filename = fileobj.name
42 else: filename = ''
Guido van Rossum68de3791997-07-19 20:22:23 +000043 if mode is None:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000044 if hasattr(fileobj, 'mode'): mode = fileobj.mode
Fred Drake9bb76d11999-04-05 18:33:40 +000045 else: mode = 'rb'
Guido van Rossum68de3791997-07-19 20:22:23 +000046
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000047 if mode[0:1] == 'r':
48 self.mode = READ
Tim Peters07e99cb2001-01-14 23:47:14 +000049 # Set flag indicating start of a new member
50 self._new_member = 1
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +000051 self.extrabuf = ""
52 self.extrasize = 0
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000053 self.filename = filename
Guido van Rossum15262191997-04-30 16:04:57 +000054
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +000055 elif mode[0:1] == 'w' or mode[0:1] == 'a':
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000056 self.mode = WRITE
57 self._init_write(filename)
58 self.compress = zlib.compressobj(compresslevel,
Tim Peters07e99cb2001-01-14 23:47:14 +000059 zlib.DEFLATED,
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000060 -zlib.MAX_WBITS,
61 zlib.DEF_MEM_LEVEL,
62 0)
63 else:
64 raise ValueError, "Mode " + mode + " not supported"
Guido van Rossum15262191997-04-30 16:04:57 +000065
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000066 self.fileobj = fileobj
Guido van Rossum15262191997-04-30 16:04:57 +000067
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000068 if self.mode == WRITE:
69 self._write_gzip_header()
Guido van Rossum15262191997-04-30 16:04:57 +000070
71 def __repr__(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000072 s = repr(self.fileobj)
73 return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
Guido van Rossum15262191997-04-30 16:04:57 +000074
75 def _init_write(self, filename):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000076 if filename[-3:] != '.gz':
77 filename = filename + '.gz'
78 self.filename = filename
79 self.crc = zlib.crc32("")
80 self.size = 0
81 self.writebuf = []
82 self.bufsize = 0
Guido van Rossum15262191997-04-30 16:04:57 +000083
84 def _write_gzip_header(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000085 self.fileobj.write('\037\213') # magic header
86 self.fileobj.write('\010') # compression method
87 fname = self.filename[:-3]
88 flags = 0
89 if fname:
90 flags = FNAME
91 self.fileobj.write(chr(flags))
Guido van Rossum95bdd0b1999-04-12 14:34:16 +000092 write32u(self.fileobj, long(time.time()))
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000093 self.fileobj.write('\002')
94 self.fileobj.write('\377')
95 if fname:
96 self.fileobj.write(fname + '\000')
Guido van Rossum15262191997-04-30 16:04:57 +000097
98 def _init_read(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000099 self.crc = zlib.crc32("")
100 self.size = 0
Guido van Rossum15262191997-04-30 16:04:57 +0000101
102 def _read_gzip_header(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000103 magic = self.fileobj.read(2)
104 if magic != '\037\213':
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000105 raise IOError, 'Not a gzipped file'
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000106 method = ord( self.fileobj.read(1) )
107 if method != 8:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000108 raise IOError, 'Unknown compression method'
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000109 flag = ord( self.fileobj.read(1) )
110 # modtime = self.fileobj.read(4)
111 # extraflag = self.fileobj.read(1)
112 # os = self.fileobj.read(1)
113 self.fileobj.read(6)
Guido van Rossum15262191997-04-30 16:04:57 +0000114
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000115 if flag & FEXTRA:
116 # Read & discard the extra field, if present
Tim Peters07e99cb2001-01-14 23:47:14 +0000117 xlen=ord(self.fileobj.read(1))
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000118 xlen=xlen+256*ord(self.fileobj.read(1))
119 self.fileobj.read(xlen)
120 if flag & FNAME:
121 # Read and discard a null-terminated string containing the filename
122 while (1):
123 s=self.fileobj.read(1)
124 if not s or s=='\000': break
125 if flag & FCOMMENT:
126 # Read and discard a null-terminated string containing a comment
127 while (1):
128 s=self.fileobj.read(1)
129 if not s or s=='\000': break
130 if flag & FHCRC:
131 self.fileobj.read(2) # Read & discard the 16-bit header CRC
Guido van Rossum15262191997-04-30 16:04:57 +0000132
133
134 def write(self,data):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000135 if self.fileobj is None:
136 raise ValueError, "write() on closed GzipFile object"
137 if len(data) > 0:
138 self.size = self.size + len(data)
139 self.crc = zlib.crc32(data, self.crc)
140 self.fileobj.write( self.compress.compress(data) )
Guido van Rossum15262191997-04-30 16:04:57 +0000141
142 def writelines(self,lines):
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000143 self.write(" ".join(lines))
Guido van Rossum15262191997-04-30 16:04:57 +0000144
Guido van Rossum56068012000-02-02 16:51:06 +0000145 def read(self, size=-1):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000146 if self.extrasize <= 0 and self.fileobj is None:
147 return ''
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000148
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000149 readsize = 1024
Guido van Rossum56068012000-02-02 16:51:06 +0000150 if size < 0: # get the whole thing
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000151 try:
152 while 1:
153 self._read(readsize)
154 readsize = readsize * 2
155 except EOFError:
156 size = self.extrasize
157 else: # just get some more of it
158 try:
159 while size > self.extrasize:
160 self._read(readsize)
161 readsize = readsize * 2
162 except EOFError:
Guido van Rossum84c6fc91998-08-03 15:41:39 +0000163 if size > self.extrasize:
164 size = self.extrasize
Tim Peters07e99cb2001-01-14 23:47:14 +0000165
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000166 chunk = self.extrabuf[:size]
167 self.extrabuf = self.extrabuf[size:]
168 self.extrasize = self.extrasize - size
Guido van Rossum15262191997-04-30 16:04:57 +0000169
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000170 return chunk
Guido van Rossum15262191997-04-30 16:04:57 +0000171
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000172 def _unread(self, buf):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000173 self.extrabuf = buf + self.extrabuf
Guido van Rossum84c6fc91998-08-03 15:41:39 +0000174 self.extrasize = len(buf) + self.extrasize
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000175
176 def _read(self, size=1024):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000177 if self.fileobj is None: raise EOFError, "Reached EOF"
Tim Peters07e99cb2001-01-14 23:47:14 +0000178
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000179 if self._new_member:
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000180 # If the _new_member flag is set, we have to
181 # jump to the next member, if there is one.
Tim Peters07e99cb2001-01-14 23:47:14 +0000182 #
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000183 # First, check if we're at the end of the file;
184 # if so, it's time to stop; no more members to read.
185 pos = self.fileobj.tell() # Save current position
186 self.fileobj.seek(0, 2) # Seek to end of file
187 if pos == self.fileobj.tell():
188 self.fileobj = None
Andrew M. Kuchling2d813e51999-09-06 16:34:51 +0000189 raise EOFError, "Reached EOF"
Tim Peters07e99cb2001-01-14 23:47:14 +0000190 else:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000191 self.fileobj.seek( pos ) # Return to original position
Tim Peters07e99cb2001-01-14 23:47:14 +0000192
193 self._init_read()
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000194 self._read_gzip_header()
195 self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
196 self._new_member = 0
Tim Peters07e99cb2001-01-14 23:47:14 +0000197
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000198 # Read a chunk of data from the file
199 buf = self.fileobj.read(size)
Tim Peters07e99cb2001-01-14 23:47:14 +0000200
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000201 # If the EOF has been reached, flush the decompression object
202 # and mark this object as finished.
Tim Peters07e99cb2001-01-14 23:47:14 +0000203
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000204 if buf == "":
205 uncompress = self.decompress.flush()
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000206 self._read_eof()
207 self.fileobj = None
208 self._add_read_data( uncompress )
209 raise EOFError, 'Reached EOF'
Tim Peters07e99cb2001-01-14 23:47:14 +0000210
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000211 uncompress = self.decompress.decompress(buf)
212 self._add_read_data( uncompress )
213
214 if self.decompress.unused_data != "":
215 # Ending case: we've come to the end of a member in the file,
216 # so seek back to the start of the unused data, finish up
217 # this member, and read a new gzip header.
218 # (The number of bytes to seek back is the length of the unused
219 # data, minus 8 because _read_eof() will rewind a further 8 bytes)
220 self.fileobj.seek( -len(self.decompress.unused_data)+8, 1)
221
222 # Check the CRC and file size, and set the flag so we read
Tim Peters07e99cb2001-01-14 23:47:14 +0000223 # a new member on the next call
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000224 self._read_eof()
Tim Peters07e99cb2001-01-14 23:47:14 +0000225 self._new_member = 1
226
227 def _add_read_data(self, data):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000228 self.crc = zlib.crc32(data, self.crc)
229 self.extrabuf = self.extrabuf + data
230 self.extrasize = self.extrasize + len(data)
231 self.size = self.size + len(data)
Guido van Rossum15262191997-04-30 16:04:57 +0000232
233 def _read_eof(self):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000234 # We've read to the end of the file, so we have to rewind in order
Tim Peters07e99cb2001-01-14 23:47:14 +0000235 # to reread the 8 bytes containing the CRC and the file size.
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000236 # We check the that the computed CRC and size of the
237 # uncompressed data matches the stored values.
238 self.fileobj.seek(-8, 1)
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000239 crc32 = read32(self.fileobj)
240 isize = read32(self.fileobj)
Guido van Rossum95bdd0b1999-04-12 14:34:16 +0000241 if crc32%0x100000000L != self.crc%0x100000000L:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000242 raise ValueError, "CRC check failed"
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000243 elif isize != self.size:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000244 raise ValueError, "Incorrect length of data produced"
Tim Peters07e99cb2001-01-14 23:47:14 +0000245
Guido van Rossum15262191997-04-30 16:04:57 +0000246 def close(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000247 if self.mode == WRITE:
248 self.fileobj.write(self.compress.flush())
249 write32(self.fileobj, self.crc)
250 write32(self.fileobj, self.size)
251 self.fileobj = None
252 elif self.mode == READ:
253 self.fileobj = None
254 if self.myfileobj:
255 self.myfileobj.close()
256 self.myfileobj = None
Guido van Rossum15262191997-04-30 16:04:57 +0000257
Andrew M. Kuchling916fcc31999-08-10 13:19:30 +0000258 def __del__(self):
Jeremy Hyltone298c302000-05-08 16:59:59 +0000259 try:
260 if (self.myfileobj is None and
261 self.fileobj is None):
262 return
263 except AttributeError:
264 return
265 self.close()
Tim Peters07e99cb2001-01-14 23:47:14 +0000266
Guido van Rossum15262191997-04-30 16:04:57 +0000267 def flush(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000268 self.fileobj.flush()
Guido van Rossum15262191997-04-30 16:04:57 +0000269
Guido van Rossum15262191997-04-30 16:04:57 +0000270 def isatty(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000271 return 0
Guido van Rossum15262191997-04-30 16:04:57 +0000272
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000273 def readline(self, size=-1):
274 if size < 0: size = sys.maxint
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000275 bufs = []
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000276 orig_size = size
277 readsize = min(100, size) # Read from the file in small chunks
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000278 while 1:
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000279 if size == 0:
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000280 return "".join(bufs) # Return resulting line
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000281
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000282 c = self.read(readsize)
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000283 i = c.find('\n')
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000284 if size is not None:
285 # We set i=size to break out of the loop under two
Tim Peters07e99cb2001-01-14 23:47:14 +0000286 # conditions: 1) there's no newline, and the chunk is
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000287 # larger than size, or 2) there is a newline, but the
288 # resulting line would be longer than 'size'.
289 if i==-1 and len(c) > size: i=size-1
290 elif size <= i: i = size -1
Guido van Rossum15262191997-04-30 16:04:57 +0000291
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000292 if i >= 0 or c == '':
293 bufs.append(c[:i+1]) # Add portion of last chunk
294 self._unread(c[i+1:]) # Push back rest of chunk
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000295 return ''.join(bufs) # Return resulting line
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000296
297 # Append chunk to list, decrease 'size',
298 bufs.append(c)
299 size = size - len(c)
300 readsize = min(size, readsize * 2)
Tim Peters07e99cb2001-01-14 23:47:14 +0000301
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000302 def readlines(self, sizehint=0):
303 # Negative numbers result in reading all the lines
304 if sizehint <= 0: sizehint = sys.maxint
305 L = []
306 while sizehint > 0:
307 line = self.readline()
308 if line == "": break
309 L.append( line )
310 sizehint = sizehint - len(line)
311
312 return L
Guido van Rossum15262191997-04-30 16:04:57 +0000313
Guido van Rossum68de3791997-07-19 20:22:23 +0000314 def writelines(self, L):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000315 for line in L:
316 self.write(line)
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000317
318
319def _test():
320 # Act like gzip; with -d, act like gunzip.
321 # The input file is not deleted, however, nor are any other gzip
322 # options or features supported.
323 import sys
324 args = sys.argv[1:]
325 decompress = args and args[0] == "-d"
326 if decompress:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000327 args = args[1:]
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000328 if not args:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000329 args = ["-"]
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000330 for arg in args:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000331 if decompress:
332 if arg == "-":
333 f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
334 g = sys.stdout
335 else:
336 if arg[-3:] != ".gz":
337 print "filename doesn't end in .gz:", `arg`
338 continue
339 f = open(arg, "rb")
340 g = __builtin__.open(arg[:-3], "wb")
341 else:
342 if arg == "-":
343 f = sys.stdin
344 g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
345 else:
346 f = __builtin__.open(arg, "rb")
347 g = open(arg + ".gz", "wb")
348 while 1:
349 chunk = f.read(1024)
350 if not chunk:
351 break
352 g.write(chunk)
353 if g is not sys.stdout:
354 g.close()
355 if f is not sys.stdin:
356 f.close()
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000357
358if __name__ == '__main__':
359 _test()