cpython/Lib/compression/_common/_streams.py

"""Internal classes used by compression modules"""

import io
import sys

BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE  # Compressed data read chunk size


class BaseStream(io.BufferedIOBase):
    """Mode-checking helper functions."""

    def _check_not_closed(self):
        if self.closed:
            raise ValueError("I/O operation on closed file")

    def _check_can_read(self):
        if not self.readable():
            raise io.UnsupportedOperation("File not open for reading")

    def _check_can_write(self):
        if not self.writable():
            raise io.UnsupportedOperation("File not open for writing")

    def _check_can_seek(self):
        if not self.readable():
            raise io.UnsupportedOperation("Seeking is only supported "
                                          "on files open for reading")
        if not self.seekable():
            raise io.UnsupportedOperation("The underlying file object "
                                          "does not support seeking")


class DecompressReader(io.RawIOBase):
    """Adapts the decompressor API to a RawIOBase reader API"""

    def readable(self):
        return True

    def __init__(self, fp, decomp_factory, trailing_error=(), *,
                 stream_magic=(), **decomp_args):
        self._fp = fp
        self._eof = False
        self._pos = 0  # Current offset in decompressed stream

        # Set to size of decompressed stream once it is known, for SEEK_END
        self._size = -1

        # Save the decompressor factory and arguments.
        # If the file contains multiple compressed streams, each
        # stream will need a separate decompressor object. A new decompressor
        # object is also needed when implementing a backwards seek().
        self._decomp_factory = decomp_factory
        self._decomp_args = decomp_args
        self._decompressor = self._decomp_factory(**self._decomp_args)

        # Exception class to catch from decompressor signifying invalid
        # trailing data to ignore
        self._trailing_error = trailing_error

        # Prefix (or tuple of prefixes) that every compressed stream starts
        # with, if the format has one.  Data starting with it is a damaged
        # stream rather than trailing data, so its error is not ignored.
        self._stream_magic = stream_magic

    def close(self):
        self._decompressor = None
        return super().close()

    def seekable(self):
        return self._fp.seekable()

    def readinto(self, b):
        with memoryview(b) as view, view.cast("B") as byte_view:
            data = self.read(len(byte_view))
            byte_view[:len(data)] = data
        return len(data)

    def read(self, size=-1):
        if size < 0:
            return self.readall()

        if not size or self._eof:
            return b""
        data = None  # Default if EOF is encountered
        # Depending on the input data, our call to the decompressor may not
        # return any data. In this case, try again after reading another block.
        while True:
            if self._decompressor.eof:
                rawblock = (self._decompressor.unused_data or
                            self._fp.read(BUFFER_SIZE))
                if not rawblock:
                    break
                # Continue to next stream.
                self._decompressor = self._decomp_factory(
                    **self._decomp_args)
                try:
                    data = self._decompressor.decompress(rawblock, size)
                except self._trailing_error:
                    if (self._stream_magic
                            and rawblock.startswith(self._stream_magic)):
                        raise  # Damaged stream, not trailing data.
                    # Trailing data isn't a valid compressed stream; ignore it.
                    break
            else:
                if self._decompressor.needs_input:
                    rawblock = self._fp.read(BUFFER_SIZE)
                    if not rawblock:
                        raise EOFError("Compressed file ended before the "
                                       "end-of-stream marker was reached")
                else:
                    rawblock = b""
                data = self._decompressor.decompress(rawblock, size)
            if data:
                break
        if not data:
            self._eof = True
            self._size = self._pos
            return b""
        self._pos += len(data)
        return data

    def readall(self):
        chunks = []
        # sys.maxsize means the max length of output buffer is unlimited,
        # so that the whole input buffer can be decompressed within one
        # .decompress() call.
        while data := self.read(sys.maxsize):
            chunks.append(data)

        return b"".join(chunks)

    # Rewind the file to the beginning of the data stream.
    def _rewind(self):
        self._fp.seek(0)
        self._eof = False
        self._pos = 0
        self._decompressor = self._decomp_factory(**self._decomp_args)

    def seek(self, offset, whence=io.SEEK_SET):
        # Recalculate offset as an absolute file position.
        if whence == io.SEEK_SET:
            pass
        elif whence == io.SEEK_CUR:
            offset = self._pos + offset
        elif whence == io.SEEK_END:
            # Seeking relative to EOF - we need to know the file's size.
            if self._size < 0:
                while self.read(io.DEFAULT_BUFFER_SIZE):
                    pass
            offset = self._size + offset
        else:
            raise ValueError("Invalid value for whence: {}".format(whence))

        # Make it so that offset is the number of bytes to skip forward.
        if offset < self._pos:
            self._rewind()
        else:
            offset -= self._pos

        # Read and discard data until we reach the desired position.
        while offset > 0:
            data = self.read(min(io.DEFAULT_BUFFER_SIZE, offset))
            if not data:
                break
            offset -= len(data)

        return self._pos

    def tell(self):
        """Return the current file position."""
        return self._pos


def _decompress_all(data, decomp_factory, trailing_error=(), *,
                    stream_magic=(), truncated_error, **decomp_args):
    """Decompress all of the concatenated streams in *data*.

    Each stream is decompressed by a new decompressor, made by calling
    *decomp_factory* with *decomp_args*.  As in DecompressReader, an error
    of class *trailing_error* in a stream after the first one ends the
    data, unless that stream starts with *stream_magic*.  A stream without
    an end-of-stream marker raises *truncated_error*.
    """
    results = []
    # The views are released on exit, so that the caller can resize a
    # bytearray argument again, also while handling an exception.
    with memoryview(data) as view, view.cast("B") as data:
        pos = 0
        # There is usually a single stream: the first chunk is all of the
        # data, to decompress it in one go.
        chunk_size = len(data)
        while True:
            decomp = decomp_factory(**decomp_args)
            start = pos
            res = []
            try:
                # Feed the stream in chunks of increasing size, so that what
                # is read past its end (and copied into unused_data) is
                # bounded by the size of the stream, not by the size of all
                # the streams that follow it.
                while not decomp.eof and pos < len(data):
                    with data[pos:pos + chunk_size] as chunk:
                        res.append(decomp.decompress(chunk))
                        pos += len(chunk)
                    chunk_size *= 2
            except trailing_error:
                if not results:
                    raise  # Error in the first stream; bail out.
                if stream_magic:
                    if isinstance(stream_magic, bytes):
                        stream_magic = (stream_magic,)
                    size = max(map(len, stream_magic))
                    with data[start:start + size] as head:
                        if head.tobytes().startswith(stream_magic):
                            raise  # Damaged stream, not trailing data.
                # Trailing data isn't a valid compressed stream; ignore it.
                break
            results.extend(res)
            if not decomp.eof:
                raise truncated_error("Compressed data ended before the "
                                      "end-of-stream marker was reached")
            pos -= len(decomp.unused_data)
            if pos == len(data):
                break
            # Later streams start small: a larger first chunk reads further
            # past the end of a short stream, a smaller one needs more calls.
            chunk_size = 1024
    return b"".join(results)