1# This Source Code Form is subject to the terms of the Mozilla Public 2# License, v. 2.0. If a copy of the MPL was not distributed with this 3# file, You can obtain one at http://mozilla.org/MPL/2.0/. 4 5import os 6import time 7import zipfile 8 9from mozbuild.util import lock_file 10 11 12class ZipFile(zipfile.ZipFile): 13 """ Class with methods to open, read, write, close, list zip files. 14 15 Subclassing zipfile.ZipFile to allow for overwriting of existing 16 entries, though only for writestr, not for write. 17 """ 18 def __init__(self, file, mode="r", compression=zipfile.ZIP_STORED, 19 lock = False): 20 if lock: 21 assert isinstance(file, basestring) 22 self.lockfile = lock_file(file + '.lck') 23 else: 24 self.lockfile = None 25 26 if mode == 'a' and lock: 27 # appending to a file which doesn't exist fails, but we can't check 28 # existence util we hold the lock 29 if (not os.path.isfile(file)) or os.path.getsize(file) == 0: 30 mode = 'w' 31 32 zipfile.ZipFile.__init__(self, file, mode, compression) 33 self._remove = [] 34 self.end = self.fp.tell() 35 self.debug = 0 36 37 def writestr(self, zinfo_or_arcname, bytes): 38 """Write contents into the archive. 39 40 The contents is the argument 'bytes', 'zinfo_or_arcname' is either 41 a ZipInfo instance or the name of the file in the archive. 42 This method is overloaded to allow overwriting existing entries. 43 """ 44 if not isinstance(zinfo_or_arcname, zipfile.ZipInfo): 45 zinfo = zipfile.ZipInfo(filename=zinfo_or_arcname, 46 date_time=time.localtime(time.time())) 47 zinfo.compress_type = self.compression 48 # Add some standard UNIX file access permissions (-rw-r--r--). 49 zinfo.external_attr = (0x81a4 & 0xFFFF) << 16L 50 else: 51 zinfo = zinfo_or_arcname 52 53 # Now to the point why we overwrote this in the first place, 54 # remember the entry numbers if we already had this entry. 55 # Optimizations: 56 # If the entry to overwrite is the last one, just reuse that. 57 # If we store uncompressed and the new content has the same size 58 # as the old, reuse the existing entry. 59 60 doSeek = False # store if we need to seek to the eof after overwriting 61 if self.NameToInfo.has_key(zinfo.filename): 62 # Find the last ZipInfo with our name. 63 # Last, because that's catching multiple overwrites 64 i = len(self.filelist) 65 while i > 0: 66 i -= 1 67 if self.filelist[i].filename == zinfo.filename: 68 break 69 zi = self.filelist[i] 70 if ((zinfo.compress_type == zipfile.ZIP_STORED 71 and zi.compress_size == len(bytes)) 72 or (i + 1) == len(self.filelist)): 73 # make sure we're allowed to write, otherwise done by writestr below 74 self._writecheck(zi) 75 # overwrite existing entry 76 self.fp.seek(zi.header_offset) 77 if (i + 1) == len(self.filelist): 78 # this is the last item in the file, just truncate 79 self.fp.truncate() 80 else: 81 # we need to move to the end of the file afterwards again 82 doSeek = True 83 # unhook the current zipinfo, the writestr of our superclass 84 # will add a new one 85 self.filelist.pop(i) 86 self.NameToInfo.pop(zinfo.filename) 87 else: 88 # Couldn't optimize, sadly, just remember the old entry for removal 89 self._remove.append(self.filelist.pop(i)) 90 zipfile.ZipFile.writestr(self, zinfo, bytes) 91 self.filelist.sort(lambda l, r: cmp(l.header_offset, r.header_offset)) 92 if doSeek: 93 self.fp.seek(self.end) 94 self.end = self.fp.tell() 95 96 def close(self): 97 """Close the file, and for mode "w" and "a" write the ending 98 records. 99 100 Overwritten to compact overwritten entries. 101 """ 102 if not self._remove: 103 # we don't have anything special to do, let's just call base 104 r = zipfile.ZipFile.close(self) 105 self.lockfile = None 106 return r 107 108 if self.fp.mode != 'r+b': 109 # adjust file mode if we originally just wrote, now we rewrite 110 self.fp.close() 111 self.fp = open(self.filename, 'r+b') 112 all = map(lambda zi: (zi, True), self.filelist) + \ 113 map(lambda zi: (zi, False), self._remove) 114 all.sort(lambda l, r: cmp(l[0].header_offset, r[0].header_offset)) 115 # empty _remove for multiple closes 116 self._remove = [] 117 118 lengths = [all[i+1][0].header_offset - all[i][0].header_offset 119 for i in xrange(len(all)-1)] 120 lengths.append(self.end - all[-1][0].header_offset) 121 to_pos = 0 122 for (zi, keep), length in zip(all, lengths): 123 if not keep: 124 continue 125 oldoff = zi.header_offset 126 # python <= 2.4 has file_offset 127 if hasattr(zi, 'file_offset'): 128 zi.file_offset = zi.file_offset + to_pos - oldoff 129 zi.header_offset = to_pos 130 self.fp.seek(oldoff) 131 content = self.fp.read(length) 132 self.fp.seek(to_pos) 133 self.fp.write(content) 134 to_pos += length 135 self.fp.truncate() 136 zipfile.ZipFile.close(self) 137 self.lockfile = None 138