developer.intra2net.com Git - pyi2ncommon/blob - src/zip_stream.py

   1 # The software in this package is distributed under the GNU General
   2 # Public License version 2 (with a special exception described below).
   3 #
   4 # A copy of GNU General Public License (GPL) is included in this distribution,
   5 # in the file COPYING.GPL.
   6 #
   7 # As a special exception, if other files instantiate templates or use macros
   8 # or inline functions from this file, or you compile this file and link it
   9 # with other works to produce a work based on this file, this file
  10 # does not by itself cause the resulting work to be covered
  11 # by the GNU General Public License.
  12 #
  13 # However the source code for this file must still be made available
  14 # in accordance with section (3) of the GNU General Public License.
  15 #
  16 # This exception does not invalidate any other reasons why a work based
  17 # on this file might be covered by the GNU General Public License.
  18 #
  19 # Copyright (c) 2016-2018 Intra2net AG <info@intra2net.com>
  20
  21 """ Streamable version of zipfile
  22
  23 Python's :py:class:`zipfile.ZipFile` can only write to seekable streams
  24 since version 3.5 and only implements adding files as wholes. This module
  25 implements class :py:class:`ZipStream` which is a subclass of ZipFile that can
  26 read from non-seekable input streams and write to non-seekable output streams.
  27
  28 Use as follows::
  29
  30     from pyi2ncommon.zip_stream import ZipStream
  31     with ZipStream(output_stream, 'w') as zip:
  32         info = zip.create_zipinfo(big_file)
  33         with open(big_file, 'rb') as input_stream:     # always read binary!
  34             zip.write_stream(input_stream, info)
  35
  36 .. codeauthor:: Intra2net AG <info@intra2net>
  37 """
  38
  39 import sys
  40 import os
  41
  42 if sys.version_info.major < 3:
  43     raise ImportError('Did not backport zipfile from python 3.5 to py2')
  44 if sys.version_info.minor >= 6:
  45     # imports for _write_stream_36
  46     import shutil
  47 else:
  48     # imports for create_zipinfo, _write_stream_35 and _get_compressor
  49     from stat import S_ISDIR
  50     import time
  51     import zlib
  52     crc32 = zlib.crc32
  53     import bz2
  54     import struct
  55 if sys.version_info.minor >= 5:
  56     from zipfile import *
  57 else:
  58     # backport of zipfile from python 3.5 to support stream output
  59     from zipfile35 import *
  60
  61 from .type_helpers import isstr
  62
  63
  64 # copied from zipfile.py
  65 ZIP64_LIMIT = (1 << 31) - 1
  66
  67
  68 def _get_compressor(compress_type):
  69     """Copied fomr zipfile.py in py3.5 (cannot legally import)"""
  70     if compress_type == ZIP_DEFLATED:
  71         return zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION,
  72                                 zlib.DEFLATED, -15)
  73     elif compress_type == ZIP_BZIP2:
  74         return bz2.BZ2Compressor()
  75     elif compress_type == ZIP_LZMA:
  76         return LZMACompressor()
  77     else:
  78         return None
  79
  80
  81 class BytesTellWrapper:
  82     """
  83     Wrapper around a write-only stream that supports tell but not seek
  84
  85     copy of zipfile._Tellable
  86     """
  87     def __init__(self, fp):
  88         self.fp = fp
  89         self.offset = 0
  90
  91     def write(self, data):
  92         n = self.fp.write(data)
  93         self.offset += n
  94         return n
  95
  96     def tell(self):
  97         return self.offset
  98
  99     def flush(self):
 100         self.fp.flush()
 101
 102     def close(self):
 103         self.fp.close()
 104
 105
 106 class ZipStream(ZipFile):
 107     """Subclass of ZipFile that supports non-seekable input and output"""
 108
 109     def __init__(self, file, *args, **kwargs):
 110         """
 111         Create ZipStream instance which is like a ZipFile plus functions
 112         create_zipinfo and write_stream.
 113
 114         ZipFile determines whether output stream can seek() and tell().
 115         Unfortunately some streams (like sys.stdout.buffer when redirecting
 116         output) seem to support these methods but only return 0 from tell.
 117
 118         :param bool force_wrap: force wrapping of output stream with
 119                                 BytesTellWrapper.
 120         """
 121         if 'force_wrap' in kwargs and kwargs['force_wrap']:
 122             if isstr(file):
 123                 raise ValueError('force_wrap only makes sense for streams')
 124             del kwargs['force_wrap']
 125             super(ZipStream, self).__init__(BytesTellWrapper(file), *args,
 126                                             **kwargs)
 127         else:
 128             super(ZipStream, self).__init__(file, *args, **kwargs)
 129
 130     def create_zipinfo(self, filename, arcname=None):
 131         """
 132         Create ZipInfo for given file
 133
 134         Optionally set arcname as name of file inside archive.
 135
 136         Adapted from zipfile.py in (ZipInfo.from_file in py3.6, ZipFile.write
 137         in py3.5)
 138         """
 139         if sys.version_info.major >= 3 and sys.version_info.minor >= 6:
 140             return ZipInfo.from_file(filename, arcname)
 141
 142         st = os.stat(filename)
 143         isdir = S_ISDIR(st.st_mode)
 144         mtime = time.localtime(st.st_mtime)
 145         date_time = mtime[0:6]
 146         # Create ZipInfo instance to store file information
 147         if arcname is None:
 148             arcname = filename
 149         arcname = os.path.normpath(os.path.splitdrive(arcname)[1])
 150         while arcname[0] in (os.sep, os.altsep):
 151             arcname = arcname[1:]
 152         if isdir:
 153             arcname += '/'
 154         zinfo = ZipInfo(arcname, date_time)
 155         zinfo.external_attr = (st.st_mode & 0xFFFF) << 16  # Unix attributes
 156         if isdir:
 157             zinfo.compress_type = ZIP_STORED
 158             zinfo.file_size = 0
 159             zinfo.external_attr |= 0x10  # MS-DOS directory flag
 160         else:
 161             zinfo.compress_type = self.compression
 162             zinfo.file_size = st.st_size
 163
 164         return zinfo
 165
 166     def write_stream(self, src, zinfo):
 167         """
 168         Add data from byte stream stream src to archive with info in ZipInfo.
 169
 170         Param zinfo must be a ZipInfo, created e.g. with
 171         :py:meth:`ZipStream.create_zipinfo`
 172
 173         Note: you cannot add directories this way (removed the corresponding
 174         code).
 175
 176         This is a shortened version of python's
 177         :py:func:`zipfile.ZipFile.write`.
 178         """
 179         if sys.version_info.major >= 3 and sys.version_info.minor >= 6:
 180             return self._write_stream_36(src, zinfo)
 181         else:
 182             return self._write_stream_35(src, zinfo)
 183
 184     def _write_stream_35(self, src, zinfo):
 185         """Implementation of _write_stream based on ZipFile.write (py 3.5)"""
 186         if not self.fp:
 187             raise RuntimeError(
 188                 "Attempt to write to ZIP archive that was already closed")
 189
 190         zinfo.flag_bits = 0x00
 191
 192         with self._lock:
 193             zinfo.header_offset = self.fp.tell()    # Start of header bytes
 194             if zinfo.compress_type == ZIP_LZMA:
 195                 # Compressed data includes an end-of-stream (EOS) marker
 196                 zinfo.flag_bits |= 0x02
 197
 198             self._writecheck(zinfo)
 199             self._didModify = True
 200
 201             cmpr = _get_compressor(zinfo.compress_type)
 202             zinfo.flag_bits |= 0x08
 203
 204             # Must overwrite CRC and sizes with correct data later
 205             zinfo.CRC = CRC = 0
 206             zinfo.compress_size = compress_size = 0
 207             # Compressed size can be larger than uncompressed size
 208             zip64 = self._allowZip64 and \
 209                 zinfo.file_size * 1.05 > ZIP64_LIMIT
 210             self.fp.write(zinfo.FileHeader(zip64))
 211             file_size = 0
 212             while 1:
 213                 buf = src.read(1024 * 8)
 214                 if not buf:
 215                     break
 216                 file_size = file_size + len(buf)
 217                 CRC = crc32(buf, CRC) & 0xffffffff
 218                 if cmpr:
 219                     buf = cmpr.compress(buf)
 220                     compress_size = compress_size + len(buf)
 221                 self.fp.write(buf)
 222             if cmpr:
 223                 buf = cmpr.flush()
 224                 compress_size = compress_size + len(buf)
 225                 self.fp.write(buf)
 226                 zinfo.compress_size = compress_size
 227             else:
 228                 zinfo.compress_size = file_size
 229             zinfo.CRC = CRC
 230             zinfo.file_size = file_size
 231
 232             # Write CRC and file sizes after the file data
 233             fmt = '<LQQ' if zip64 else '<LLL'
 234             self.fp.write(struct.pack(fmt, zinfo.CRC, zinfo.compress_size,
 235                                       zinfo.file_size))
 236             self.start_dir = self.fp.tell()
 237             self.filelist.append(zinfo)
 238             self.NameToInfo[zinfo.filename] = zinfo
 239
 240     def _write_stream_36(self, src, zinfo):
 241         """Implementation of _write_stream based on ZipFile.write (py 3.6)"""
 242         if not self.fp:
 243             raise ValueError(
 244                 "Attempt to write to ZIP archive that was already closed")
 245         if self._writing:
 246             raise ValueError(
 247                 "Can't write to ZIP archive while an open writing handle exists"
 248             )
 249
 250         if zinfo.is_dir():
 251             raise ValueError('streaming a dir entry does not make sense')
 252         if zinfo.compress_type is None:
 253             zinfo.compress_type = self.compression
 254
 255         with self.open(zinfo, 'w') as dest:
 256             shutil.copyfileobj(src, dest, 1024*8)