michael@0: # This Source Code Form is subject to the terms of the Mozilla Public michael@0: # License, v. 2.0. If a copy of the MPL was not distributed with this michael@0: # file, You can obtain one at http://mozilla.org/MPL/2.0/. michael@0: michael@0: import os michael@0: import time michael@0: import zipfile michael@0: michael@0: from mozbuild.util import lock_file michael@0: michael@0: michael@0: class ZipFile(zipfile.ZipFile): michael@0: """ Class with methods to open, read, write, close, list zip files. michael@0: michael@0: Subclassing zipfile.ZipFile to allow for overwriting of existing michael@0: entries, though only for writestr, not for write. michael@0: """ michael@0: def __init__(self, file, mode="r", compression=zipfile.ZIP_STORED, michael@0: lock = False): michael@0: if lock: michael@0: assert isinstance(file, basestring) michael@0: self.lockfile = lock_file(file + '.lck') michael@0: else: michael@0: self.lockfile = None michael@0: michael@0: if mode == 'a' and lock: michael@0: # appending to a file which doesn't exist fails, but we can't check michael@0: # existence util we hold the lock michael@0: if (not os.path.isfile(file)) or os.path.getsize(file) == 0: michael@0: mode = 'w' michael@0: michael@0: zipfile.ZipFile.__init__(self, file, mode, compression) michael@0: self._remove = [] michael@0: self.end = self.fp.tell() michael@0: self.debug = 0 michael@0: michael@0: def writestr(self, zinfo_or_arcname, bytes): michael@0: """Write contents into the archive. michael@0: michael@0: The contents is the argument 'bytes', 'zinfo_or_arcname' is either michael@0: a ZipInfo instance or the name of the file in the archive. michael@0: This method is overloaded to allow overwriting existing entries. michael@0: """ michael@0: if not isinstance(zinfo_or_arcname, zipfile.ZipInfo): michael@0: zinfo = zipfile.ZipInfo(filename=zinfo_or_arcname, michael@0: date_time=time.localtime(time.time())) michael@0: zinfo.compress_type = self.compression michael@0: # Add some standard UNIX file access permissions (-rw-r--r--). michael@0: zinfo.external_attr = (0x81a4 & 0xFFFF) << 16L michael@0: else: michael@0: zinfo = zinfo_or_arcname michael@0: michael@0: # Now to the point why we overwrote this in the first place, michael@0: # remember the entry numbers if we already had this entry. michael@0: # Optimizations: michael@0: # If the entry to overwrite is the last one, just reuse that. michael@0: # If we store uncompressed and the new content has the same size michael@0: # as the old, reuse the existing entry. michael@0: michael@0: doSeek = False # store if we need to seek to the eof after overwriting michael@0: if self.NameToInfo.has_key(zinfo.filename): michael@0: # Find the last ZipInfo with our name. michael@0: # Last, because that's catching multiple overwrites michael@0: i = len(self.filelist) michael@0: while i > 0: michael@0: i -= 1 michael@0: if self.filelist[i].filename == zinfo.filename: michael@0: break michael@0: zi = self.filelist[i] michael@0: if ((zinfo.compress_type == zipfile.ZIP_STORED michael@0: and zi.compress_size == len(bytes)) michael@0: or (i + 1) == len(self.filelist)): michael@0: # make sure we're allowed to write, otherwise done by writestr below michael@0: self._writecheck(zi) michael@0: # overwrite existing entry michael@0: self.fp.seek(zi.header_offset) michael@0: if (i + 1) == len(self.filelist): michael@0: # this is the last item in the file, just truncate michael@0: self.fp.truncate() michael@0: else: michael@0: # we need to move to the end of the file afterwards again michael@0: doSeek = True michael@0: # unhook the current zipinfo, the writestr of our superclass michael@0: # will add a new one michael@0: self.filelist.pop(i) michael@0: self.NameToInfo.pop(zinfo.filename) michael@0: else: michael@0: # Couldn't optimize, sadly, just remember the old entry for removal michael@0: self._remove.append(self.filelist.pop(i)) michael@0: zipfile.ZipFile.writestr(self, zinfo, bytes) michael@0: self.filelist.sort(lambda l, r: cmp(l.header_offset, r.header_offset)) michael@0: if doSeek: michael@0: self.fp.seek(self.end) michael@0: self.end = self.fp.tell() michael@0: michael@0: def close(self): michael@0: """Close the file, and for mode "w" and "a" write the ending michael@0: records. michael@0: michael@0: Overwritten to compact overwritten entries. michael@0: """ michael@0: if not self._remove: michael@0: # we don't have anything special to do, let's just call base michael@0: r = zipfile.ZipFile.close(self) michael@0: self.lockfile = None michael@0: return r michael@0: michael@0: if self.fp.mode != 'r+b': michael@0: # adjust file mode if we originally just wrote, now we rewrite michael@0: self.fp.close() michael@0: self.fp = open(self.filename, 'r+b') michael@0: all = map(lambda zi: (zi, True), self.filelist) + \ michael@0: map(lambda zi: (zi, False), self._remove) michael@0: all.sort(lambda l, r: cmp(l[0].header_offset, r[0].header_offset)) michael@0: # empty _remove for multiple closes michael@0: self._remove = [] michael@0: michael@0: lengths = [all[i+1][0].header_offset - all[i][0].header_offset michael@0: for i in xrange(len(all)-1)] michael@0: lengths.append(self.end - all[-1][0].header_offset) michael@0: to_pos = 0 michael@0: for (zi, keep), length in zip(all, lengths): michael@0: if not keep: michael@0: continue michael@0: oldoff = zi.header_offset michael@0: # python <= 2.4 has file_offset michael@0: if hasattr(zi, 'file_offset'): michael@0: zi.file_offset = zi.file_offset + to_pos - oldoff michael@0: zi.header_offset = to_pos michael@0: self.fp.seek(oldoff) michael@0: content = self.fp.read(length) michael@0: self.fp.seek(to_pos) michael@0: self.fp.write(content) michael@0: to_pos += length michael@0: self.fp.truncate() michael@0: zipfile.ZipFile.close(self) michael@0: self.lockfile = None