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