config/MozZipFile.py

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