1 files changed, 143 insertions, 0 deletions
diff --git a/config/MozZipFile.py b/config/MozZipFile.py
new file mode 100644
index 0000000000..6e74bc1eba
--- /dev/null
+++ b/config/MozZipFile.py
@@ -0,0 +1,143 @@
+# 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
+
+import six
+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, six.text_type)
+            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) << 16
+        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 zinfo.filename in self.NameToInfo:
+            # 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(key=lambda l: l.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(key=lambda l: l[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