1 files changed, 137 insertions, 0 deletions
diff --git a/config/MozZipFile.py b/config/MozZipFile.py
new file mode 100644
index 000000000..337fe0521
--- /dev/null
+++ b/config/MozZipFile.py
@@ -0,0 +1,137 @@
+# 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