blob: 9df8e5d180ee6f81b4d56faf2b93bfe559bb7f12 [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
Martin v. Löwis8cc965c2001-08-09 07:21:56 +000067 self.offset = 0
Guido van Rossum15262191997-04-30 16:04:57 +000068
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000069 if self.mode == WRITE:
70 self._write_gzip_header()
Guido van Rossum15262191997-04-30 16:04:57 +000071
72 def __repr__(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000073 s = repr(self.fileobj)
74 return '<gzip ' + s[1:-1] + ' ' + hex(id(self)) + '>'
Guido van Rossum15262191997-04-30 16:04:57 +000075
76 def _init_write(self, filename):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000077 if filename[-3:] != '.gz':
78 filename = filename + '.gz'
79 self.filename = filename
80 self.crc = zlib.crc32("")
81 self.size = 0
82 self.writebuf = []
83 self.bufsize = 0
Guido van Rossum15262191997-04-30 16:04:57 +000084
85 def _write_gzip_header(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000086 self.fileobj.write('\037\213') # magic header
87 self.fileobj.write('\010') # compression method
88 fname = self.filename[:-3]
89 flags = 0
90 if fname:
91 flags = FNAME
92 self.fileobj.write(chr(flags))
Guido van Rossum95bdd0b1999-04-12 14:34:16 +000093 write32u(self.fileobj, long(time.time()))
Guido van Rossum45e2fbc1998-03-26 21:13:24 +000094 self.fileobj.write('\002')
95 self.fileobj.write('\377')
96 if fname:
97 self.fileobj.write(fname + '\000')
Guido van Rossum15262191997-04-30 16:04:57 +000098
99 def _init_read(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000100 self.crc = zlib.crc32("")
101 self.size = 0
Guido van Rossum15262191997-04-30 16:04:57 +0000102
103 def _read_gzip_header(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000104 magic = self.fileobj.read(2)
105 if magic != '\037\213':
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000106 raise IOError, 'Not a gzipped file'
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000107 method = ord( self.fileobj.read(1) )
108 if method != 8:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000109 raise IOError, 'Unknown compression method'
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000110 flag = ord( self.fileobj.read(1) )
111 # modtime = self.fileobj.read(4)
112 # extraflag = self.fileobj.read(1)
113 # os = self.fileobj.read(1)
114 self.fileobj.read(6)
Guido van Rossum15262191997-04-30 16:04:57 +0000115
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000116 if flag & FEXTRA:
117 # Read & discard the extra field, if present
Tim Peters07e99cb2001-01-14 23:47:14 +0000118 xlen=ord(self.fileobj.read(1))
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000119 xlen=xlen+256*ord(self.fileobj.read(1))
120 self.fileobj.read(xlen)
121 if flag & FNAME:
122 # Read and discard a null-terminated string containing the filename
123 while (1):
124 s=self.fileobj.read(1)
125 if not s or s=='\000': break
126 if flag & FCOMMENT:
127 # Read and discard a null-terminated string containing a comment
128 while (1):
129 s=self.fileobj.read(1)
130 if not s or s=='\000': break
131 if flag & FHCRC:
132 self.fileobj.read(2) # Read & discard the 16-bit header CRC
Guido van Rossum15262191997-04-30 16:04:57 +0000133
134
135 def write(self,data):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000136 if self.fileobj is None:
137 raise ValueError, "write() on closed GzipFile object"
138 if len(data) > 0:
139 self.size = self.size + len(data)
140 self.crc = zlib.crc32(data, self.crc)
141 self.fileobj.write( self.compress.compress(data) )
Martin v. Löwis8cc965c2001-08-09 07:21:56 +0000142 self.offset += len(data)
Guido van Rossum15262191997-04-30 16:04:57 +0000143
144 def writelines(self,lines):
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000145 self.write(" ".join(lines))
Guido van Rossum15262191997-04-30 16:04:57 +0000146
Guido van Rossum56068012000-02-02 16:51:06 +0000147 def read(self, size=-1):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000148 if self.extrasize <= 0 and self.fileobj is None:
149 return ''
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000150
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000151 readsize = 1024
Guido van Rossum56068012000-02-02 16:51:06 +0000152 if size < 0: # get the whole thing
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000153 try:
154 while 1:
155 self._read(readsize)
156 readsize = readsize * 2
157 except EOFError:
158 size = self.extrasize
159 else: # just get some more of it
160 try:
161 while size > self.extrasize:
162 self._read(readsize)
163 readsize = readsize * 2
164 except EOFError:
Guido van Rossum84c6fc91998-08-03 15:41:39 +0000165 if size > self.extrasize:
166 size = self.extrasize
Tim Peters07e99cb2001-01-14 23:47:14 +0000167
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000168 chunk = self.extrabuf[:size]
169 self.extrabuf = self.extrabuf[size:]
170 self.extrasize = self.extrasize - size
Guido van Rossum15262191997-04-30 16:04:57 +0000171
Martin v. Löwis8cc965c2001-08-09 07:21:56 +0000172 self.offset += size
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000173 return chunk
Guido van Rossum15262191997-04-30 16:04:57 +0000174
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000175 def _unread(self, buf):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000176 self.extrabuf = buf + self.extrabuf
Guido van Rossum84c6fc91998-08-03 15:41:39 +0000177 self.extrasize = len(buf) + self.extrasize
Martin v. Löwis8cc965c2001-08-09 07:21:56 +0000178 self.offset -= len(buf)
Guido van Rossumb16a3b81998-01-27 19:29:45 +0000179
180 def _read(self, size=1024):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000181 if self.fileobj is None: raise EOFError, "Reached EOF"
Tim Peters07e99cb2001-01-14 23:47:14 +0000182
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000183 if self._new_member:
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000184 # If the _new_member flag is set, we have to
185 # jump to the next member, if there is one.
Tim Peters07e99cb2001-01-14 23:47:14 +0000186 #
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000187 # First, check if we're at the end of the file;
188 # if so, it's time to stop; no more members to read.
189 pos = self.fileobj.tell() # Save current position
190 self.fileobj.seek(0, 2) # Seek to end of file
191 if pos == self.fileobj.tell():
Andrew M. Kuchling2d813e51999-09-06 16:34:51 +0000192 raise EOFError, "Reached EOF"
Tim Peters07e99cb2001-01-14 23:47:14 +0000193 else:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000194 self.fileobj.seek( pos ) # Return to original position
Tim Peters07e99cb2001-01-14 23:47:14 +0000195
196 self._init_read()
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000197 self._read_gzip_header()
198 self.decompress = zlib.decompressobj(-zlib.MAX_WBITS)
199 self._new_member = 0
Tim Peters07e99cb2001-01-14 23:47:14 +0000200
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000201 # Read a chunk of data from the file
202 buf = self.fileobj.read(size)
Tim Peters07e99cb2001-01-14 23:47:14 +0000203
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000204 # If the EOF has been reached, flush the decompression object
205 # and mark this object as finished.
Tim Peters07e99cb2001-01-14 23:47:14 +0000206
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000207 if buf == "":
208 uncompress = self.decompress.flush()
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000209 self._read_eof()
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000210 self._add_read_data( uncompress )
211 raise EOFError, 'Reached EOF'
Tim Peters07e99cb2001-01-14 23:47:14 +0000212
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000213 uncompress = self.decompress.decompress(buf)
214 self._add_read_data( uncompress )
215
216 if self.decompress.unused_data != "":
217 # Ending case: we've come to the end of a member in the file,
218 # so seek back to the start of the unused data, finish up
219 # this member, and read a new gzip header.
220 # (The number of bytes to seek back is the length of the unused
221 # data, minus 8 because _read_eof() will rewind a further 8 bytes)
222 self.fileobj.seek( -len(self.decompress.unused_data)+8, 1)
223
224 # Check the CRC and file size, and set the flag so we read
Tim Peters07e99cb2001-01-14 23:47:14 +0000225 # a new member on the next call
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000226 self._read_eof()
Tim Peters07e99cb2001-01-14 23:47:14 +0000227 self._new_member = 1
228
229 def _add_read_data(self, data):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000230 self.crc = zlib.crc32(data, self.crc)
231 self.extrabuf = self.extrabuf + data
232 self.extrasize = self.extrasize + len(data)
233 self.size = self.size + len(data)
Guido van Rossum15262191997-04-30 16:04:57 +0000234
235 def _read_eof(self):
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000236 # We've read to the end of the file, so we have to rewind in order
Tim Peters07e99cb2001-01-14 23:47:14 +0000237 # to reread the 8 bytes containing the CRC and the file size.
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000238 # We check the that the computed CRC and size of the
239 # uncompressed data matches the stored values.
240 self.fileobj.seek(-8, 1)
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000241 crc32 = read32(self.fileobj)
242 isize = read32(self.fileobj)
Guido van Rossum95bdd0b1999-04-12 14:34:16 +0000243 if crc32%0x100000000L != self.crc%0x100000000L:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000244 raise ValueError, "CRC check failed"
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000245 elif isize != self.size:
Andrew M. Kuchlingf4f119c1999-03-25 21:49:14 +0000246 raise ValueError, "Incorrect length of data produced"
Tim Peters07e99cb2001-01-14 23:47:14 +0000247
Guido van Rossum15262191997-04-30 16:04:57 +0000248 def close(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000249 if self.mode == WRITE:
250 self.fileobj.write(self.compress.flush())
251 write32(self.fileobj, self.crc)
252 write32(self.fileobj, self.size)
253 self.fileobj = None
254 elif self.mode == READ:
255 self.fileobj = None
256 if self.myfileobj:
257 self.myfileobj.close()
258 self.myfileobj = None
Guido van Rossum15262191997-04-30 16:04:57 +0000259
Andrew M. Kuchling916fcc31999-08-10 13:19:30 +0000260 def __del__(self):
Jeremy Hyltone298c302000-05-08 16:59:59 +0000261 try:
262 if (self.myfileobj is None and
263 self.fileobj is None):
264 return
265 except AttributeError:
266 return
267 self.close()
Tim Peters07e99cb2001-01-14 23:47:14 +0000268
Guido van Rossum15262191997-04-30 16:04:57 +0000269 def flush(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000270 self.fileobj.flush()
Guido van Rossum15262191997-04-30 16:04:57 +0000271
Guido van Rossum15262191997-04-30 16:04:57 +0000272 def isatty(self):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000273 return 0
Guido van Rossum15262191997-04-30 16:04:57 +0000274
Martin v. Löwis8cc965c2001-08-09 07:21:56 +0000275 def tell(self):
276 return self.offset
277
278 def rewind(self):
279 '''Return the uncompressed stream file position indicator to the
280 beginning of the file'''
281 if self.mode != READ:
282 raise IOError("Can't rewind in write mode")
283 self.fileobj.seek(0)
284 self._new_member = 1
285 self.extrabuf = ""
286 self.extrasize = 0
287 self.offset = 0
288
289 def seek(self, offset):
290 if self.mode == WRITE:
291 if offset < self.offset:
292 raise IOError('Negative seek in write mode')
293 count = offset - self.offset
294 for i in range(count/1024):
295 f.write(1024*'\0')
296 self.write((count%1024)*'\0')
297 elif self.mode == READ:
298 if offset < self.offset:
299 # for negative seek, rewind and do positive seek
300 self.rewind()
301 count = offset - self.offset
302 for i in range(count/1024): self.read(1024)
303 self.read(count % 1024)
304
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000305 def readline(self, size=-1):
306 if size < 0: size = sys.maxint
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000307 bufs = []
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000308 orig_size = size
309 readsize = min(100, size) # Read from the file in small chunks
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000310 while 1:
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000311 if size == 0:
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000312 return "".join(bufs) # Return resulting line
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000313
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000314 c = self.read(readsize)
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000315 i = c.find('\n')
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000316 if size is not None:
317 # We set i=size to break out of the loop under two
Tim Peters07e99cb2001-01-14 23:47:14 +0000318 # conditions: 1) there's no newline, and the chunk is
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000319 # larger than size, or 2) there is a newline, but the
320 # resulting line would be longer than 'size'.
321 if i==-1 and len(c) > size: i=size-1
322 elif size <= i: i = size -1
Guido van Rossum15262191997-04-30 16:04:57 +0000323
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000324 if i >= 0 or c == '':
325 bufs.append(c[:i+1]) # Add portion of last chunk
326 self._unread(c[i+1:]) # Push back rest of chunk
Eric S. Raymondee5e61d2001-02-09 09:10:35 +0000327 return ''.join(bufs) # Return resulting line
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000328
329 # Append chunk to list, decrease 'size',
330 bufs.append(c)
331 size = size - len(c)
332 readsize = min(size, readsize * 2)
Tim Peters07e99cb2001-01-14 23:47:14 +0000333
Andrew M. Kuchling41616ee2000-07-29 20:15:26 +0000334 def readlines(self, sizehint=0):
335 # Negative numbers result in reading all the lines
336 if sizehint <= 0: sizehint = sys.maxint
337 L = []
338 while sizehint > 0:
339 line = self.readline()
340 if line == "": break
341 L.append( line )
342 sizehint = sizehint - len(line)
343
344 return L
Guido van Rossum15262191997-04-30 16:04:57 +0000345
Guido van Rossum68de3791997-07-19 20:22:23 +0000346 def writelines(self, L):
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000347 for line in L:
348 self.write(line)
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000349
350
351def _test():
352 # Act like gzip; with -d, act like gunzip.
353 # The input file is not deleted, however, nor are any other gzip
354 # options or features supported.
355 import sys
356 args = sys.argv[1:]
357 decompress = args and args[0] == "-d"
358 if decompress:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000359 args = args[1:]
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000360 if not args:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000361 args = ["-"]
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000362 for arg in args:
Guido van Rossum45e2fbc1998-03-26 21:13:24 +0000363 if decompress:
364 if arg == "-":
365 f = GzipFile(filename="", mode="rb", fileobj=sys.stdin)
366 g = sys.stdout
367 else:
368 if arg[-3:] != ".gz":
369 print "filename doesn't end in .gz:", `arg`
370 continue
371 f = open(arg, "rb")
372 g = __builtin__.open(arg[:-3], "wb")
373 else:
374 if arg == "-":
375 f = sys.stdin
376 g = GzipFile(filename="", mode="wb", fileobj=sys.stdout)
377 else:
378 f = __builtin__.open(arg, "rb")
379 g = open(arg + ".gz", "wb")
380 while 1:
381 chunk = f.read(1024)
382 if not chunk:
383 break
384 g.write(chunk)
385 if g is not sys.stdout:
386 g.close()
387 if f is not sys.stdin:
388 f.close()
Guido van Rossum51ca6e31997-12-30 20:09:08 +0000389
390if __name__ == '__main__':
391 _test()