Clean up, remove compat with py < 3.6

[pyi2ncommon] / src / zip_stream.py

# The software in this package is distributed under the GNU General
# Public License version 2 (with a special exception described below).
#
# A copy of GNU General Public License (GPL) is included in this distribution,
# in the file COPYING.GPL.
#
# As a special exception, if other files instantiate templates or use macros
# or inline functions from this file, or you compile this file and link it
# with other works to produce a work based on this file, this file
# does not by itself cause the resulting work to be covered
# by the GNU General Public License.
#
# However the source code for this file must still be made available
# in accordance with section (3) of the GNU General Public License.
#
# This exception does not invalidate any other reasons why a work based
# on this file might be covered by the GNU General Public License.
#
# Copyright (c) 2016-2018 Intra2net AG <info@intra2net.com>

""" Streamable version of zipfile

Python's :py:class:`zipfile.ZipFile` can only write to seekable streams
since version 3.5 and only implements adding files as wholes. This module
implements class :py:class:`ZipStream` which is a subclass of ZipFile that can
read from non-seekable input streams and write to non-seekable output streams.

Use as follows::

    from pyi2ncommon.zip_stream import ZipStream
    with ZipStream(output_stream, 'w') as zip:
        info = zip.create_zipinfo(big_file)
        with open(big_file, 'rb') as input_stream:     # always read binary!
            zip.write_stream(input_stream, info)

.. codeauthor:: Intra2net AG <info@intra2net>
"""

import shutil
from zipfile import *

from .type_helpers import isstr


class BytesTellWrapper:
    """
    Wrapper around a write-only stream that supports tell but not seek

    copy of zipfile._Tellable
    """
    def __init__(self, fp):
        self.fp = fp
        self.offset = 0

    def write(self, data):
        n = self.fp.write(data)
        self.offset += n
        return n

    def tell(self):
        return self.offset

    def flush(self):
        self.fp.flush()

    def close(self):
        self.fp.close()


class ZipStream(ZipFile):
    """Subclass of ZipFile that supports non-seekable input and output"""

    def __init__(self, file, *args, **kwargs):
        """
        Create ZipStream instance which is like a ZipFile plus functions
        create_zipinfo and write_stream.

        ZipFile determines whether output stream can seek() and tell().
        Unfortunately some streams (like sys.stdout.buffer when redirecting
        output) seem to support these methods but only return 0 from tell.

        :param bool force_wrap: force wrapping of output stream with
                                BytesTellWrapper.
        """
        if 'force_wrap' in kwargs and kwargs['force_wrap']:
            if isstr(file):
                raise ValueError('force_wrap only makes sense for streams')
            del kwargs['force_wrap']
            super(ZipStream, self).__init__(BytesTellWrapper(file), *args,
                                            **kwargs)
        else:
            super(ZipStream, self).__init__(file, *args, **kwargs)

    def create_zipinfo(self, filename, arcname=None):
        """
        Create ZipInfo for given file

        Optionally set arcname as name of file inside archive.
        """
        return ZipInfo.from_file(filename, arcname)

    def write_stream(self, src, zinfo):
        """
        Add data from byte stream stream src to archive with info in ZipInfo.

        Param zinfo must be a ZipInfo, created e.g. with
        :py:meth:`ZipStream.create_zipinfo`

        Note: you cannot add directories this way (removed the corresponding
        code).

        This is a shortened version of python's
        :py:func:`zipfile.ZipFile.write`.
        """
        if not self.fp:
            raise ValueError(
                "Attempt to write to ZIP archive that was already closed")
        if self._writing:
            raise ValueError(
                "Can't write to ZIP archive while an open writing handle exists"
            )

        if zinfo.is_dir():
            raise ValueError('streaming a dir entry does not make sense')
        if zinfo.compress_type is None:
            zinfo.compress_type = self.compression

        with self.open(zinfo, 'w') as dest:
            shutil.copyfileobj(src, dest, 1024*8)