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 shutil
  40 from zipfile import *
  41
  42 from .type_helpers import isstr
  43
  44
  45 class BytesTellWrapper:
  46     """
  47     Wrapper around a write-only stream that supports tell but not seek
  48
  49     copy of zipfile._Tellable
  50     """
  51     def __init__(self, fp):
  52         self.fp = fp
  53         self.offset = 0
  54
  55     def write(self, data):
  56         n = self.fp.write(data)
  57         self.offset += n
  58         return n
  59
  60     def tell(self):
  61         return self.offset
  62
  63     def flush(self):
  64         self.fp.flush()
  65
  66     def close(self):
  67         self.fp.close()
  68
  69
  70 class ZipStream(ZipFile):
  71     """Subclass of ZipFile that supports non-seekable input and output"""
  72
  73     def __init__(self, file, *args, **kwargs):
  74         """
  75         Create ZipStream instance which is like a ZipFile plus functions
  76         create_zipinfo and write_stream.
  77
  78         ZipFile determines whether output stream can seek() and tell().
  79         Unfortunately some streams (like sys.stdout.buffer when redirecting
  80         output) seem to support these methods but only return 0 from tell.
  81
  82         :param bool force_wrap: force wrapping of output stream with
  83                                 BytesTellWrapper.
  84         """
  85         if 'force_wrap' in kwargs and kwargs['force_wrap']:
  86             if isstr(file):
  87                 raise ValueError('force_wrap only makes sense for streams')
  88             del kwargs['force_wrap']
  89             super(ZipStream, self).__init__(BytesTellWrapper(file), *args,
  90                                             **kwargs)
  91         else:
  92             super(ZipStream, self).__init__(file, *args, **kwargs)
  93
  94     def create_zipinfo(self, filename, arcname=None):
  95         """
  96         Create ZipInfo for given file
  97
  98         Optionally set arcname as name of file inside archive.
  99         """
 100         return ZipInfo.from_file(filename, arcname)
 101
 102     def write_stream(self, src, zinfo):
 103         """
 104         Add data from byte stream stream src to archive with info in ZipInfo.
 105
 106         Param zinfo must be a ZipInfo, created e.g. with
 107         :py:meth:`ZipStream.create_zipinfo`
 108
 109         Note: you cannot add directories this way (removed the corresponding
 110         code).
 111
 112         This is a shortened version of python's
 113         :py:func:`zipfile.ZipFile.write`.
 114         """
 115         if not self.fp:
 116             raise ValueError(
 117                 "Attempt to write to ZIP archive that was already closed")
 118         if self._writing:
 119             raise ValueError(
 120                 "Can't write to ZIP archive while an open writing handle exists"
 121             )
 122
 123         if zinfo.is_dir():
 124             raise ValueError('streaming a dir entry does not make sense')
 125         if zinfo.compress_type is None:
 126             zinfo.compress_type = self.compression
 127
 128         with self.open(zinfo, 'w') as dest:
 129             shutil.copyfileobj(src, dest, 1024*8)