source: trunk/essentials/dev-lang/python/Lib/zipfile.py

Last change on this file was 3225, checked in by bird, 18 years ago

Python 2.5
File size: 34.8 KB
Line 
1"""
2Read and write ZIP files.
3"""
4import struct, os, time, sys
5import binascii, cStringIO
6
7try:
8 import zlib # We may need its compression method
9except ImportError:
10 zlib = None
11
12__all__ = ["BadZipfile", "error", "ZIP_STORED", "ZIP_DEFLATED", "is_zipfile",
13 "ZipInfo", "ZipFile", "PyZipFile", "LargeZipFile" ]
14
15class BadZipfile(Exception):
16 pass
17
18
19class LargeZipFile(Exception):
20 """
21 Raised when writing a zipfile, the zipfile requires ZIP64 extensions
22 and those extensions are disabled.
23 """
24
25error = BadZipfile # The exception raised by this module
26
27ZIP64_LIMIT= (1 << 31) - 1
28
29# constants for Zip file compression methods
30ZIP_STORED = 0
31ZIP_DEFLATED = 8
32# Other ZIP compression methods not supported
33
34# Here are some struct module formats for reading headers
35structEndArchive = "<4s4H2lH" # 9 items, end of archive, 22 bytes
36stringEndArchive = "PK\005\006" # magic number for end of archive record
37structCentralDir = "<4s4B4HlLL5HLl"# 19 items, central directory, 46 bytes
38stringCentralDir = "PK\001\002" # magic number for central directory
39structFileHeader = "<4s2B4HlLL2H" # 12 items, file header record, 30 bytes
40stringFileHeader = "PK\003\004" # magic number for file header
41structEndArchive64Locator = "<4slql" # 4 items, locate Zip64 header, 20 bytes
42stringEndArchive64Locator = "PK\x06\x07" # magic token for locator header
43structEndArchive64 = "<4sqhhllqqqq" # 10 items, end of archive (Zip64), 56 bytes
44stringEndArchive64 = "PK\x06\x06" # magic token for Zip64 header
45
46
47# indexes of entries in the central directory structure
48_CD_SIGNATURE = 0
49_CD_CREATE_VERSION = 1
50_CD_CREATE_SYSTEM = 2
51_CD_EXTRACT_VERSION = 3
52_CD_EXTRACT_SYSTEM = 4 # is this meaningful?
53_CD_FLAG_BITS = 5
54_CD_COMPRESS_TYPE = 6
55_CD_TIME = 7
56_CD_DATE = 8
57_CD_CRC = 9
58_CD_COMPRESSED_SIZE = 10
59_CD_UNCOMPRESSED_SIZE = 11
60_CD_FILENAME_LENGTH = 12
61_CD_EXTRA_FIELD_LENGTH = 13
62_CD_COMMENT_LENGTH = 14
63_CD_DISK_NUMBER_START = 15
64_CD_INTERNAL_FILE_ATTRIBUTES = 16
65_CD_EXTERNAL_FILE_ATTRIBUTES = 17
66_CD_LOCAL_HEADER_OFFSET = 18
67
68# indexes of entries in the local file header structure
69_FH_SIGNATURE = 0
70_FH_EXTRACT_VERSION = 1
71_FH_EXTRACT_SYSTEM = 2 # is this meaningful?
72_FH_GENERAL_PURPOSE_FLAG_BITS = 3
73_FH_COMPRESSION_METHOD = 4
74_FH_LAST_MOD_TIME = 5
75_FH_LAST_MOD_DATE = 6
76_FH_CRC = 7
77_FH_COMPRESSED_SIZE = 8
78_FH_UNCOMPRESSED_SIZE = 9
79_FH_FILENAME_LENGTH = 10
80_FH_EXTRA_FIELD_LENGTH = 11
81
82def is_zipfile(filename):
83 """Quickly see if file is a ZIP file by checking the magic number."""
84 try:
85 fpin = open(filename, "rb")
86 endrec = _EndRecData(fpin)
87 fpin.close()
88 if endrec:
89 return True # file has correct magic number
90 except IOError:
91 pass
92 return False
93
94def _EndRecData64(fpin, offset, endrec):
95 """
96 Read the ZIP64 end-of-archive records and use that to update endrec
97 """
98 locatorSize = struct.calcsize(structEndArchive64Locator)
99 fpin.seek(offset - locatorSize, 2)
100 data = fpin.read(locatorSize)
101 sig, diskno, reloff, disks = struct.unpack(structEndArchive64Locator, data)
102 if sig != stringEndArchive64Locator:
103 return endrec
104
105 if diskno != 0 or disks != 1:
106 raise BadZipfile("zipfiles that span multiple disks are not supported")
107
108 # Assume no 'zip64 extensible data'
109 endArchiveSize = struct.calcsize(structEndArchive64)
110 fpin.seek(offset - locatorSize - endArchiveSize, 2)
111 data = fpin.read(endArchiveSize)
112 sig, sz, create_version, read_version, disk_num, disk_dir, \
113 dircount, dircount2, dirsize, diroffset = \
114 struct.unpack(structEndArchive64, data)
115 if sig != stringEndArchive64:
116 return endrec
117
118 # Update the original endrec using data from the ZIP64 record
119 endrec[1] = disk_num
120 endrec[2] = disk_dir
121 endrec[3] = dircount
122 endrec[4] = dircount2
123 endrec[5] = dirsize
124 endrec[6] = diroffset
125 return endrec
126
127
128def _EndRecData(fpin):
129 """Return data from the "End of Central Directory" record, or None.
130
131 The data is a list of the nine items in the ZIP "End of central dir"
132 record followed by a tenth item, the file seek offset of this record."""
133 fpin.seek(-22, 2) # Assume no archive comment.
134 filesize = fpin.tell() + 22 # Get file size
135 data = fpin.read()
136 if data[0:4] == stringEndArchive and data[-2:] == "\000\000":
137 endrec = struct.unpack(structEndArchive, data)
138 endrec = list(endrec)
139 endrec.append("") # Append the archive comment
140 endrec.append(filesize - 22) # Append the record start offset
141 if endrec[-4] == -1 or endrec[-4] == 0xffffffff:
142 return _EndRecData64(fpin, -22, endrec)
143 return endrec
144 # Search the last END_BLOCK bytes of the file for the record signature.
145 # The comment is appended to the ZIP file and has a 16 bit length.
146 # So the comment may be up to 64K long. We limit the search for the
147 # signature to a few Kbytes at the end of the file for efficiency.
148 # also, the signature must not appear in the comment.
149 END_BLOCK = min(filesize, 1024 * 4)
150 fpin.seek(filesize - END_BLOCK, 0)
151 data = fpin.read()
152 start = data.rfind(stringEndArchive)
153 if start >= 0: # Correct signature string was found
154 endrec = struct.unpack(structEndArchive, data[start:start+22])
155 endrec = list(endrec)
156 comment = data[start+22:]
157 if endrec[7] == len(comment): # Comment length checks out
158 # Append the archive comment and start offset
159 endrec.append(comment)
160 endrec.append(filesize - END_BLOCK + start)
161 if endrec[-4] == -1 or endrec[-4] == 0xffffffff:
162 return _EndRecData64(fpin, - END_BLOCK + start, endrec)
163 return endrec
164 return # Error, return None
165
166
167class ZipInfo (object):
168 """Class with attributes describing each file in the ZIP archive."""
169
170 __slots__ = (
171 'orig_filename',
172 'filename',
173 'date_time',
174 'compress_type',
175 'comment',
176 'extra',
177 'create_system',
178 'create_version',
179 'extract_version',
180 'reserved',
181 'flag_bits',
182 'volume',
183 'internal_attr',
184 'external_attr',
185 'header_offset',
186 'CRC',
187 'compress_size',
188 'file_size',
189 )
190
191 def __init__(self, filename="NoName", date_time=(1980,1,1,0,0,0)):
192 self.orig_filename = filename # Original file name in archive
193
194 # Terminate the file name at the first null byte. Null bytes in file
195 # names are used as tricks by viruses in archives.
196 null_byte = filename.find(chr(0))
197 if null_byte >= 0:
198 filename = filename[0:null_byte]
199 # This is used to ensure paths in generated ZIP files always use
200 # forward slashes as the directory separator, as required by the
201 # ZIP format specification.
202 if os.sep != "/" and os.sep in filename:
203 filename = filename.replace(os.sep, "/")
204
205 self.filename = filename # Normalized file name
206 self.date_time = date_time # year, month, day, hour, min, sec
207 # Standard values:
208 self.compress_type = ZIP_STORED # Type of compression for the file
209 self.comment = "" # Comment for each file
210 self.extra = "" # ZIP extra data
211 if sys.platform == 'win32':
212 self.create_system = 0 # System which created ZIP archive
213 else:
214 # Assume everything else is unix-y
215 self.create_system = 3 # System which created ZIP archive
216 self.create_version = 20 # Version which created ZIP archive
217 self.extract_version = 20 # Version needed to extract archive
218 self.reserved = 0 # Must be zero
219 self.flag_bits = 0 # ZIP flag bits
220 self.volume = 0 # Volume number of file header
221 self.internal_attr = 0 # Internal attributes
222 self.external_attr = 0 # External file attributes
223 # Other attributes are set by class ZipFile:
224 # header_offset Byte offset to the file header
225 # CRC CRC-32 of the uncompressed file
226 # compress_size Size of the compressed file
227 # file_size Size of the uncompressed file
228
229 def FileHeader(self):
230 """Return the per-file header as a string."""
231 dt = self.date_time
232 dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
233 dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
234 if self.flag_bits & 0x08:
235 # Set these to zero because we write them after the file data
236 CRC = compress_size = file_size = 0
237 else:
238 CRC = self.CRC
239 compress_size = self.compress_size
240 file_size = self.file_size
241
242 extra = self.extra
243
244 if file_size > ZIP64_LIMIT or compress_size > ZIP64_LIMIT:
245 # File is larger than what fits into a 4 byte integer,
246 # fall back to the ZIP64 extension
247 fmt = '<hhqq'
248 extra = extra + struct.pack(fmt,
249 1, struct.calcsize(fmt)-4, file_size, compress_size)
250 file_size = 0xffffffff # -1
251 compress_size = 0xffffffff # -1
252 self.extract_version = max(45, self.extract_version)
253 self.create_version = max(45, self.extract_version)
254
255 header = struct.pack(structFileHeader, stringFileHeader,
256 self.extract_version, self.reserved, self.flag_bits,
257 self.compress_type, dostime, dosdate, CRC,
258 compress_size, file_size,
259 len(self.filename), len(extra))
260 return header + self.filename + extra
261
262 def _decodeExtra(self):
263 # Try to decode the extra field.
264 extra = self.extra
265 unpack = struct.unpack
266 while extra:
267 tp, ln = unpack('<hh', extra[:4])
268 if tp == 1:
269 if ln >= 24:
270 counts = unpack('<qqq', extra[4:28])
271 elif ln == 16:
272 counts = unpack('<qq', extra[4:20])
273 elif ln == 8:
274 counts = unpack('<q', extra[4:12])
275 elif ln == 0:
276 counts = ()
277 else:
278 raise RuntimeError, "Corrupt extra field %s"%(ln,)
279
280 idx = 0
281
282 # ZIP64 extension (large files and/or large archives)
283 if self.file_size == -1 or self.file_size == 0xFFFFFFFFL:
284 self.file_size = counts[idx]
285 idx += 1
286
287 if self.compress_size == -1 or self.compress_size == 0xFFFFFFFFL:
288 self.compress_size = counts[idx]
289 idx += 1
290
291 if self.header_offset == -1 or self.header_offset == 0xffffffffL:
292 old = self.header_offset
293 self.header_offset = counts[idx]
294 idx+=1
295
296 extra = extra[ln+4:]
297
298
299class ZipFile:
300 """ Class with methods to open, read, write, close, list zip files.
301
302 z = ZipFile(file, mode="r", compression=ZIP_STORED, allowZip64=True)
303
304 file: Either the path to the file, or a file-like object.
305 If it is a path, the file will be opened and closed by ZipFile.
306 mode: The mode can be either read "r", write "w" or append "a".
307 compression: ZIP_STORED (no compression) or ZIP_DEFLATED (requires zlib).
308 allowZip64: if True ZipFile will create files with ZIP64 extensions when
309 needed, otherwise it will raise an exception when this would
310 be necessary.
311
312 """
313
314 fp = None # Set here since __del__ checks it
315
316 def __init__(self, file, mode="r", compression=ZIP_STORED, allowZip64=False):
317 """Open the ZIP file with mode read "r", write "w" or append "a"."""
318 self._allowZip64 = allowZip64
319 self._didModify = False
320 if compression == ZIP_STORED:
321 pass
322 elif compression == ZIP_DEFLATED:
323 if not zlib:
324 raise RuntimeError,\
325 "Compression requires the (missing) zlib module"
326 else:
327 raise RuntimeError, "That compression method is not supported"
328 self.debug = 0 # Level of printing: 0 through 3
329 self.NameToInfo = {} # Find file info given name
330 self.filelist = [] # List of ZipInfo instances for archive
331 self.compression = compression # Method of compression
332 self.mode = key = mode.replace('b', '')[0]
333
334 # Check if we were passed a file-like object
335 if isinstance(file, basestring):
336 self._filePassed = 0
337 self.filename = file
338 modeDict = {'r' : 'rb', 'w': 'wb', 'a' : 'r+b'}
339 self.fp = open(file, modeDict[mode])
340 else:
341 self._filePassed = 1
342 self.fp = file
343 self.filename = getattr(file, 'name', None)
344
345 if key == 'r':
346 self._GetContents()
347 elif key == 'w':
348 pass
349 elif key == 'a':
350 try: # See if file is a zip file
351 self._RealGetContents()
352 # seek to start of directory and overwrite
353 self.fp.seek(self.start_dir, 0)
354 except BadZipfile: # file is not a zip file, just append
355 self.fp.seek(0, 2)
356 else:
357 if not self._filePassed:
358 self.fp.close()
359 self.fp = None
360 raise RuntimeError, 'Mode must be "r", "w" or "a"'
361
362 def _GetContents(self):
363 """Read the directory, making sure we close the file if the format
364 is bad."""
365 try:
366 self._RealGetContents()
367 except BadZipfile:
368 if not self._filePassed:
369 self.fp.close()
370 self.fp = None
371 raise
372
373 def _RealGetContents(self):
374 """Read in the table of contents for the ZIP file."""
375 fp = self.fp
376 endrec = _EndRecData(fp)
377 if not endrec:
378 raise BadZipfile, "File is not a zip file"
379 if self.debug > 1:
380 print endrec
381 size_cd = endrec[5] # bytes in central directory
382 offset_cd = endrec[6] # offset of central directory
383 self.comment = endrec[8] # archive comment
384 # endrec[9] is the offset of the "End of Central Dir" record
385 if endrec[9] > ZIP64_LIMIT:
386 x = endrec[9] - size_cd - 56 - 20
387 else:
388 x = endrec[9] - size_cd
389 # "concat" is zero, unless zip was concatenated to another file
390 concat = x - offset_cd
391 if self.debug > 2:
392 print "given, inferred, offset", offset_cd, x, concat
393 # self.start_dir: Position of start of central directory
394 self.start_dir = offset_cd + concat
395 fp.seek(self.start_dir, 0)
396 data = fp.read(size_cd)
397 fp = cStringIO.StringIO(data)
398 total = 0
399 while total < size_cd:
400 centdir = fp.read(46)
401 total = total + 46
402 if centdir[0:4] != stringCentralDir:
403 raise BadZipfile, "Bad magic number for central directory"
404 centdir = struct.unpack(structCentralDir, centdir)
405 if self.debug > 2:
406 print centdir
407 filename = fp.read(centdir[_CD_FILENAME_LENGTH])
408 # Create ZipInfo instance to store file information
409 x = ZipInfo(filename)
410 x.extra = fp.read(centdir[_CD_EXTRA_FIELD_LENGTH])
411 x.comment = fp.read(centdir[_CD_COMMENT_LENGTH])
412 total = (total + centdir[_CD_FILENAME_LENGTH]
413 + centdir[_CD_EXTRA_FIELD_LENGTH]
414 + centdir[_CD_COMMENT_LENGTH])
415 x.header_offset = centdir[_CD_LOCAL_HEADER_OFFSET]
416 (x.create_version, x.create_system, x.extract_version, x.reserved,
417 x.flag_bits, x.compress_type, t, d,
418 x.CRC, x.compress_size, x.file_size) = centdir[1:12]
419 x.volume, x.internal_attr, x.external_attr = centdir[15:18]
420 # Convert date/time code to (year, month, day, hour, min, sec)
421 x.date_time = ( (d>>9)+1980, (d>>5)&0xF, d&0x1F,
422 t>>11, (t>>5)&0x3F, (t&0x1F) * 2 )
423
424 x._decodeExtra()
425 x.header_offset = x.header_offset + concat
426 self.filelist.append(x)
427 self.NameToInfo[x.filename] = x
428 if self.debug > 2:
429 print "total", total
430
431
432 def namelist(self):
433 """Return a list of file names in the archive."""
434 l = []
435 for data in self.filelist:
436 l.append(data.filename)
437 return l
438
439 def infolist(self):
440 """Return a list of class ZipInfo instances for files in the
441 archive."""
442 return self.filelist
443
444 def printdir(self):
445 """Print a table of contents for the zip file."""
446 print "%-46s %19s %12s" % ("File Name", "Modified ", "Size")
447 for zinfo in self.filelist:
448 date = "%d-%02d-%02d %02d:%02d:%02d" % zinfo.date_time
449 print "%-46s %s %12d" % (zinfo.filename, date, zinfo.file_size)
450
451 def testzip(self):
452 """Read all the files and check the CRC."""
453 for zinfo in self.filelist:
454 try:
455 self.read(zinfo.filename) # Check CRC-32
456 except BadZipfile:
457 return zinfo.filename
458
459
460 def getinfo(self, name):
461 """Return the instance of ZipInfo given 'name'."""
462 return self.NameToInfo[name]
463
464 def read(self, name):
465 """Return file bytes (as a string) for name."""
466 if self.mode not in ("r", "a"):
467 raise RuntimeError, 'read() requires mode "r" or "a"'
468 if not self.fp:
469 raise RuntimeError, \
470 "Attempt to read ZIP archive that was already closed"
471 zinfo = self.getinfo(name)
472 filepos = self.fp.tell()
473
474 self.fp.seek(zinfo.header_offset, 0)
475
476 # Skip the file header:
477 fheader = self.fp.read(30)
478 if fheader[0:4] != stringFileHeader:
479 raise BadZipfile, "Bad magic number for file header"
480
481 fheader = struct.unpack(structFileHeader, fheader)
482 fname = self.fp.read(fheader[_FH_FILENAME_LENGTH])
483 if fheader[_FH_EXTRA_FIELD_LENGTH]:
484 self.fp.read(fheader[_FH_EXTRA_FIELD_LENGTH])
485
486 if fname != zinfo.orig_filename:
487 raise BadZipfile, \
488 'File name in directory "%s" and header "%s" differ.' % (
489 zinfo.orig_filename, fname)
490
491 bytes = self.fp.read(zinfo.compress_size)
492 self.fp.seek(filepos, 0)
493 if zinfo.compress_type == ZIP_STORED:
494 pass
495 elif zinfo.compress_type == ZIP_DEFLATED:
496 if not zlib:
497 raise RuntimeError, \
498 "De-compression requires the (missing) zlib module"
499 # zlib compress/decompress code by Jeremy Hylton of CNRI
500 dc = zlib.decompressobj(-15)
501 bytes = dc.decompress(bytes)
502 # need to feed in unused pad byte so that zlib won't choke
503 ex = dc.decompress('Z') + dc.flush()
504 if ex:
505 bytes = bytes + ex
506 else:
507 raise BadZipfile, \
508 "Unsupported compression method %d for file %s" % \
509 (zinfo.compress_type, name)
510 crc = binascii.crc32(bytes)
511 if crc != zinfo.CRC:
512 raise BadZipfile, "Bad CRC-32 for file %s" % name
513 return bytes
514
515 def _writecheck(self, zinfo):
516 """Check for errors before writing a file to the archive."""
517 if zinfo.filename in self.NameToInfo:
518 if self.debug: # Warning for duplicate names
519 print "Duplicate name:", zinfo.filename
520 if self.mode not in ("w", "a"):
521 raise RuntimeError, 'write() requires mode "w" or "a"'
522 if not self.fp:
523 raise RuntimeError, \
524 "Attempt to write ZIP archive that was already closed"
525 if zinfo.compress_type == ZIP_DEFLATED and not zlib:
526 raise RuntimeError, \
527 "Compression requires the (missing) zlib module"
528 if zinfo.compress_type not in (ZIP_STORED, ZIP_DEFLATED):
529 raise RuntimeError, \
530 "That compression method is not supported"
531 if zinfo.file_size > ZIP64_LIMIT:
532 if not self._allowZip64:
533 raise LargeZipFile("Filesize would require ZIP64 extensions")
534 if zinfo.header_offset > ZIP64_LIMIT:
535 if not self._allowZip64:
536 raise LargeZipFile("Zipfile size would require ZIP64 extensions")
537
538 def write(self, filename, arcname=None, compress_type=None):
539 """Put the bytes from filename into the archive under the name
540 arcname."""
541 st = os.stat(filename)
542 mtime = time.localtime(st.st_mtime)
543 date_time = mtime[0:6]
544 # Create ZipInfo instance to store file information
545 if arcname is None:
546 arcname = filename
547 arcname = os.path.normpath(os.path.splitdrive(arcname)[1])
548 while arcname[0] in (os.sep, os.altsep):
549 arcname = arcname[1:]
550 zinfo = ZipInfo(arcname, date_time)
551 zinfo.external_attr = (st[0] & 0xFFFF) << 16L # Unix attributes
552 if compress_type is None:
553 zinfo.compress_type = self.compression
554 else:
555 zinfo.compress_type = compress_type
556
557 zinfo.file_size = st.st_size
558 zinfo.flag_bits = 0x00
559 zinfo.header_offset = self.fp.tell() # Start of header bytes
560
561 self._writecheck(zinfo)
562 self._didModify = True
563 fp = open(filename, "rb")
564 # Must overwrite CRC and sizes with correct data later
565 zinfo.CRC = CRC = 0
566 zinfo.compress_size = compress_size = 0
567 zinfo.file_size = file_size = 0
568 self.fp.write(zinfo.FileHeader())
569 if zinfo.compress_type == ZIP_DEFLATED:
570 cmpr = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
571 zlib.DEFLATED, -15)
572 else:
573 cmpr = None
574 while 1:
575 buf = fp.read(1024 * 8)
576 if not buf:
577 break
578 file_size = file_size + len(buf)
579 CRC = binascii.crc32(buf, CRC)
580 if cmpr:
581 buf = cmpr.compress(buf)
582 compress_size = compress_size + len(buf)
583 self.fp.write(buf)
584 fp.close()
585 if cmpr:
586 buf = cmpr.flush()
587 compress_size = compress_size + len(buf)
588 self.fp.write(buf)
589 zinfo.compress_size = compress_size
590 else:
591 zinfo.compress_size = file_size
592 zinfo.CRC = CRC
593 zinfo.file_size = file_size
594 # Seek backwards and write CRC and file sizes
595 position = self.fp.tell() # Preserve current position in file
596 self.fp.seek(zinfo.header_offset + 14, 0)
597 self.fp.write(struct.pack("<lLL", zinfo.CRC, zinfo.compress_size,
598 zinfo.file_size))
599 self.fp.seek(position, 0)
600 self.filelist.append(zinfo)
601 self.NameToInfo[zinfo.filename] = zinfo
602
603 def writestr(self, zinfo_or_arcname, bytes):
604 """Write a file into the archive. The contents is the string
605 'bytes'. 'zinfo_or_arcname' is either a ZipInfo instance or
606 the name of the file in the archive."""
607 if not isinstance(zinfo_or_arcname, ZipInfo):
608 zinfo = ZipInfo(filename=zinfo_or_arcname,
609 date_time=time.localtime(time.time()))
610 zinfo.compress_type = self.compression
611 else:
612 zinfo = zinfo_or_arcname
613 zinfo.file_size = len(bytes) # Uncompressed size
614 zinfo.header_offset = self.fp.tell() # Start of header bytes
615 self._writecheck(zinfo)
616 self._didModify = True
617 zinfo.CRC = binascii.crc32(bytes) # CRC-32 checksum
618 if zinfo.compress_type == ZIP_DEFLATED:
619 co = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
620 zlib.DEFLATED, -15)
621 bytes = co.compress(bytes) + co.flush()
622 zinfo.compress_size = len(bytes) # Compressed size
623 else:
624 zinfo.compress_size = zinfo.file_size
625 zinfo.header_offset = self.fp.tell() # Start of header bytes
626 self.fp.write(zinfo.FileHeader())
627 self.fp.write(bytes)
628 self.fp.flush()
629 if zinfo.flag_bits & 0x08:
630 # Write CRC and file sizes after the file data
631 self.fp.write(struct.pack("<lLL", zinfo.CRC, zinfo.compress_size,
632 zinfo.file_size))
633 self.filelist.append(zinfo)
634 self.NameToInfo[zinfo.filename] = zinfo
635
636 def __del__(self):
637 """Call the "close()" method in case the user forgot."""
638 self.close()
639
640 def close(self):
641 """Close the file, and for mode "w" and "a" write the ending
642 records."""
643 if self.fp is None:
644 return
645
646 if self.mode in ("w", "a") and self._didModify: # write ending records
647 count = 0
648 pos1 = self.fp.tell()
649 for zinfo in self.filelist: # write central directory
650 count = count + 1
651 dt = zinfo.date_time
652 dosdate = (dt[0] - 1980) << 9 | dt[1] << 5 | dt[2]
653 dostime = dt[3] << 11 | dt[4] << 5 | (dt[5] // 2)
654 extra = []
655 if zinfo.file_size > ZIP64_LIMIT \
656 or zinfo.compress_size > ZIP64_LIMIT:
657 extra.append(zinfo.file_size)
658 extra.append(zinfo.compress_size)
659 file_size = 0xffffffff #-1
660 compress_size = 0xffffffff #-1
661 else:
662 file_size = zinfo.file_size
663 compress_size = zinfo.compress_size
664
665 if zinfo.header_offset > ZIP64_LIMIT:
666 extra.append(zinfo.header_offset)
667 header_offset = -1 # struct "l" format: 32 one bits
668 else:
669 header_offset = zinfo.header_offset
670
671 extra_data = zinfo.extra
672 if extra:
673 # Append a ZIP64 field to the extra's
674 extra_data = struct.pack(
675 '<hh' + 'q'*len(extra),
676 1, 8*len(extra), *extra) + extra_data
677
678 extract_version = max(45, zinfo.extract_version)
679 create_version = max(45, zinfo.create_version)
680 else:
681 extract_version = zinfo.extract_version
682 create_version = zinfo.create_version
683
684 centdir = struct.pack(structCentralDir,
685 stringCentralDir, create_version,
686 zinfo.create_system, extract_version, zinfo.reserved,
687 zinfo.flag_bits, zinfo.compress_type, dostime, dosdate,
688 zinfo.CRC, compress_size, file_size,
689 len(zinfo.filename), len(extra_data), len(zinfo.comment),
690 0, zinfo.internal_attr, zinfo.external_attr,
691 header_offset)
692 self.fp.write(centdir)
693 self.fp.write(zinfo.filename)
694 self.fp.write(extra_data)
695 self.fp.write(zinfo.comment)
696
697 pos2 = self.fp.tell()
698 # Write end-of-zip-archive record
699 if pos1 > ZIP64_LIMIT:
700 # Need to write the ZIP64 end-of-archive records
701 zip64endrec = struct.pack(
702 structEndArchive64, stringEndArchive64,
703 44, 45, 45, 0, 0, count, count, pos2 - pos1, pos1)
704 self.fp.write(zip64endrec)
705
706 zip64locrec = struct.pack(
707 structEndArchive64Locator,
708 stringEndArchive64Locator, 0, pos2, 1)
709 self.fp.write(zip64locrec)
710
711 # XXX Why is `pos3` computed next? It's never referenced.
712 pos3 = self.fp.tell()
713 endrec = struct.pack(structEndArchive, stringEndArchive,
714 0, 0, count, count, pos2 - pos1, -1, 0)
715 self.fp.write(endrec)
716
717 else:
718 endrec = struct.pack(structEndArchive, stringEndArchive,
719 0, 0, count, count, pos2 - pos1, pos1, 0)
720 self.fp.write(endrec)
721 self.fp.flush()
722 if not self._filePassed:
723 self.fp.close()
724 self.fp = None
725
726
727class PyZipFile(ZipFile):
728 """Class to create ZIP archives with Python library files and packages."""
729
730 def writepy(self, pathname, basename = ""):
731 """Add all files from "pathname" to the ZIP archive.
732
733 If pathname is a package directory, search the directory and
734 all package subdirectories recursively for all *.py and enter
735 the modules into the archive. If pathname is a plain
736 directory, listdir *.py and enter all modules. Else, pathname
737 must be a Python *.py file and the module will be put into the
738 archive. Added modules are always module.pyo or module.pyc.
739 This method will compile the module.py into module.pyc if
740 necessary.
741 """
742 dir, name = os.path.split(pathname)
743 if os.path.isdir(pathname):
744 initname = os.path.join(pathname, "__init__.py")
745 if os.path.isfile(initname):
746 # This is a package directory, add it
747 if basename:
748 basename = "%s/%s" % (basename, name)
749 else:
750 basename = name
751 if self.debug:
752 print "Adding package in", pathname, "as", basename
753 fname, arcname = self._get_codename(initname[0:-3], basename)
754 if self.debug:
755 print "Adding", arcname
756 self.write(fname, arcname)
757 dirlist = os.listdir(pathname)
758 dirlist.remove("__init__.py")
759 # Add all *.py files and package subdirectories
760 for filename in dirlist:
761 path = os.path.join(pathname, filename)
762 root, ext = os.path.splitext(filename)
763 if os.path.isdir(path):
764 if os.path.isfile(os.path.join(path, "__init__.py")):
765 # This is a package directory, add it
766 self.writepy(path, basename) # Recursive call
767 elif ext == ".py":
768 fname, arcname = self._get_codename(path[0:-3],
769 basename)
770 if self.debug:
771 print "Adding", arcname
772 self.write(fname, arcname)
773 else:
774 # This is NOT a package directory, add its files at top level
775 if self.debug:
776 print "Adding files from directory", pathname
777 for filename in os.listdir(pathname):
778 path = os.path.join(pathname, filename)
779 root, ext = os.path.splitext(filename)
780 if ext == ".py":
781 fname, arcname = self._get_codename(path[0:-3],
782 basename)
783 if self.debug:
784 print "Adding", arcname
785 self.write(fname, arcname)
786 else:
787 if pathname[-3:] != ".py":
788 raise RuntimeError, \
789 'Files added with writepy() must end with ".py"'
790 fname, arcname = self._get_codename(pathname[0:-3], basename)
791 if self.debug:
792 print "Adding file", arcname
793 self.write(fname, arcname)
794
795 def _get_codename(self, pathname, basename):
796 """Return (filename, archivename) for the path.
797
798 Given a module name path, return the correct file path and
799 archive name, compiling if necessary. For example, given
800 /python/lib/string, return (/python/lib/string.pyc, string).
801 """
802 file_py = pathname + ".py"
803 file_pyc = pathname + ".pyc"
804 file_pyo = pathname + ".pyo"
805 if os.path.isfile(file_pyo) and \
806 os.stat(file_pyo).st_mtime >= os.stat(file_py).st_mtime:
807 fname = file_pyo # Use .pyo file
808 elif not os.path.isfile(file_pyc) or \
809 os.stat(file_pyc).st_mtime < os.stat(file_py).st_mtime:
810 import py_compile
811 if self.debug:
812 print "Compiling", file_py
813 try:
814 py_compile.compile(file_py, file_pyc, None, True)
815 except py_compile.PyCompileError,err:
816 print err.msg
817 fname = file_pyc
818 else:
819 fname = file_pyc
820 archivename = os.path.split(fname)[1]
821 if basename:
822 archivename = "%s/%s" % (basename, archivename)
823 return (fname, archivename)
824
825
826def main(args = None):
827 import textwrap
828 USAGE=textwrap.dedent("""\
829 Usage:
830 zipfile.py -l zipfile.zip # Show listing of a zipfile
831 zipfile.py -t zipfile.zip # Test if a zipfile is valid
832 zipfile.py -e zipfile.zip target # Extract zipfile into target dir
833 zipfile.py -c zipfile.zip src ... # Create zipfile from sources
834 """)
835 if args is None:
836 args = sys.argv[1:]
837
838 if not args or args[0] not in ('-l', '-c', '-e', '-t'):
839 print USAGE
840 sys.exit(1)
841
842 if args[0] == '-l':
843 if len(args) != 2:
844 print USAGE
845 sys.exit(1)
846 zf = ZipFile(args[1], 'r')
847 zf.printdir()
848 zf.close()
849
850 elif args[0] == '-t':
851 if len(args) != 2:
852 print USAGE
853 sys.exit(1)
854 zf = ZipFile(args[1], 'r')
855 zf.testzip()
856 print "Done testing"
857
858 elif args[0] == '-e':
859 if len(args) != 3:
860 print USAGE
861 sys.exit(1)
862
863 zf = ZipFile(args[1], 'r')
864 out = args[2]
865 for path in zf.namelist():
866 if path.startswith('./'):
867 tgt = os.path.join(out, path[2:])
868 else:
869 tgt = os.path.join(out, path)
870
871 tgtdir = os.path.dirname(tgt)
872 if not os.path.exists(tgtdir):
873 os.makedirs(tgtdir)
874 fp = open(tgt, 'wb')
875 fp.write(zf.read(path))
876 fp.close()
877 zf.close()
878
879 elif args[0] == '-c':
880 if len(args) < 3:
881 print USAGE
882 sys.exit(1)
883
884 def addToZip(zf, path, zippath):
885 if os.path.isfile(path):
886 zf.write(path, zippath, ZIP_DEFLATED)
887 elif os.path.isdir(path):
888 for nm in os.listdir(path):
889 addToZip(zf,
890 os.path.join(path, nm), os.path.join(zippath, nm))
891 # else: ignore
892
893 zf = ZipFile(args[1], 'w', allowZip64=True)
894 for src in args[2:]:
895 addToZip(zf, src, os.path.basename(src))
896
897 zf.close()
898
899if __name__ == "__main__":
900 main()
Note: See TracBrowser for help on using the repository browser.