File size: 19,281 Bytes

7885a28

"""Classes and functions for managing compressors."""

import io
import zlib

from joblib.backports import LooseVersion

try:
    from threading import RLock
except ImportError:
    from dummy_threading import RLock

try:
    import bz2
except ImportError:
    bz2 = None

try:
    import lz4
    from lz4.frame import LZ4FrameFile
except ImportError:
    lz4 = None

try:
    import lzma
except ImportError:
    lzma = None


LZ4_NOT_INSTALLED_ERROR = (
    "LZ4 is not installed. Install it with pip: https://python-lz4.readthedocs.io/"
)

# Registered compressors
_COMPRESSORS = {}

# Magic numbers of supported compression file formats.
_ZFILE_PREFIX = b"ZF"  # used with pickle files created before 0.9.3.
_ZLIB_PREFIX = b"\x78"
_GZIP_PREFIX = b"\x1f\x8b"
_BZ2_PREFIX = b"BZ"
_XZ_PREFIX = b"\xfd\x37\x7a\x58\x5a"
_LZMA_PREFIX = b"\x5d\x00"
_LZ4_PREFIX = b"\x04\x22\x4d\x18"


def register_compressor(compressor_name, compressor, force=False):
    """Register a new compressor.

    Parameters
    ----------
    compressor_name: str.
        The name of the compressor.
    compressor: CompressorWrapper
        An instance of a 'CompressorWrapper'.
    """
    global _COMPRESSORS
    if not isinstance(compressor_name, str):
        raise ValueError(
            "Compressor name should be a string, '{}' given.".format(compressor_name)
        )

    if not isinstance(compressor, CompressorWrapper):
        raise ValueError(
            "Compressor should implement the CompressorWrapper "
            "interface, '{}' given.".format(compressor)
        )

    if compressor.fileobj_factory is not None and (
        not hasattr(compressor.fileobj_factory, "read")
        or not hasattr(compressor.fileobj_factory, "write")
        or not hasattr(compressor.fileobj_factory, "seek")
        or not hasattr(compressor.fileobj_factory, "tell")
    ):
        raise ValueError(
            "Compressor 'fileobj_factory' attribute should "
            "implement the file object interface, '{}' given.".format(
                compressor.fileobj_factory
            )
        )

    if compressor_name in _COMPRESSORS and not force:
        raise ValueError("Compressor '{}' already registered.".format(compressor_name))

    _COMPRESSORS[compressor_name] = compressor


class CompressorWrapper:
    """A wrapper around a compressor file object.

    Attributes
    ----------
    obj: a file-like object
        The object must implement the buffer interface and will be used
        internally to compress/decompress the data.
    prefix: bytestring
        A bytestring corresponding to the magic number that identifies the
        file format associated to the compressor.
    extension: str
        The file extension used to automatically select this compressor during
        a dump to a file.
    """

    def __init__(self, obj, prefix=b"", extension=""):
        self.fileobj_factory = obj
        self.prefix = prefix
        self.extension = extension

    def compressor_file(self, fileobj, compresslevel=None):
        """Returns an instance of a compressor file object."""
        if compresslevel is None:
            return self.fileobj_factory(fileobj, "wb")
        else:
            return self.fileobj_factory(fileobj, "wb", compresslevel=compresslevel)

    def decompressor_file(self, fileobj):
        """Returns an instance of a decompressor file object."""
        return self.fileobj_factory(fileobj, "rb")


class BZ2CompressorWrapper(CompressorWrapper):
    prefix = _BZ2_PREFIX
    extension = ".bz2"

    def __init__(self):
        if bz2 is not None:
            self.fileobj_factory = bz2.BZ2File
        else:
            self.fileobj_factory = None

    def _check_versions(self):
        if bz2 is None:
            raise ValueError(
                "bz2 module is not compiled on your python standard library."
            )

    def compressor_file(self, fileobj, compresslevel=None):
        """Returns an instance of a compressor file object."""
        self._check_versions()
        if compresslevel is None:
            return self.fileobj_factory(fileobj, "wb")
        else:
            return self.fileobj_factory(fileobj, "wb", compresslevel=compresslevel)

    def decompressor_file(self, fileobj):
        """Returns an instance of a decompressor file object."""
        self._check_versions()
        fileobj = self.fileobj_factory(fileobj, "rb")
        return fileobj


class LZMACompressorWrapper(CompressorWrapper):
    prefix = _LZMA_PREFIX
    extension = ".lzma"
    _lzma_format_name = "FORMAT_ALONE"

    def __init__(self):
        if lzma is not None:
            self.fileobj_factory = lzma.LZMAFile
            self._lzma_format = getattr(lzma, self._lzma_format_name)
        else:
            self.fileobj_factory = None

    def _check_versions(self):
        if lzma is None:
            raise ValueError(
                "lzma module is not compiled on your python standard library."
            )

    def compressor_file(self, fileobj, compresslevel=None):
        """Returns an instance of a compressor file object."""
        if compresslevel is None:
            return self.fileobj_factory(fileobj, "wb", format=self._lzma_format)
        else:
            return self.fileobj_factory(
                fileobj, "wb", format=self._lzma_format, preset=compresslevel
            )

    def decompressor_file(self, fileobj):
        """Returns an instance of a decompressor file object."""
        return lzma.LZMAFile(fileobj, "rb")


class XZCompressorWrapper(LZMACompressorWrapper):
    prefix = _XZ_PREFIX
    extension = ".xz"
    _lzma_format_name = "FORMAT_XZ"


class LZ4CompressorWrapper(CompressorWrapper):
    prefix = _LZ4_PREFIX
    extension = ".lz4"

    def __init__(self):
        if lz4 is not None:
            self.fileobj_factory = LZ4FrameFile
        else:
            self.fileobj_factory = None

    def _check_versions(self):
        if lz4 is None:
            raise ValueError(LZ4_NOT_INSTALLED_ERROR)
        lz4_version = lz4.__version__
        if lz4_version.startswith("v"):
            lz4_version = lz4_version[1:]
        if LooseVersion(lz4_version) < LooseVersion("0.19"):
            raise ValueError(LZ4_NOT_INSTALLED_ERROR)

    def compressor_file(self, fileobj, compresslevel=None):
        """Returns an instance of a compressor file object."""
        self._check_versions()
        if compresslevel is None:
            return self.fileobj_factory(fileobj, "wb")
        else:
            return self.fileobj_factory(fileobj, "wb", compression_level=compresslevel)

    def decompressor_file(self, fileobj):
        """Returns an instance of a decompressor file object."""
        self._check_versions()
        return self.fileobj_factory(fileobj, "rb")


###############################################################################
#  base file compression/decompression object definition
_MODE_CLOSED = 0
_MODE_READ = 1
_MODE_READ_EOF = 2
_MODE_WRITE = 3
_BUFFER_SIZE = 8192


class BinaryZlibFile(io.BufferedIOBase):
    """A file object providing transparent zlib (de)compression.

    TODO python2_drop: is it still needed since we dropped Python 2 support A
    BinaryZlibFile can act as a wrapper for an existing file object, or refer
    directly to a named file on disk.

    Note that BinaryZlibFile provides only a *binary* file interface: data read
    is returned as bytes, and data to be written should be given as bytes.

    This object is an adaptation of the BZ2File object and is compatible with
    versions of python >= 2.7.

    If filename is a str or bytes object, it gives the name
    of the file to be opened. Otherwise, it should be a file object,
    which will be used to read or write the compressed data.

    mode can be 'rb' for reading (default) or 'wb' for (over)writing

    If mode is 'wb', compresslevel can be a number between 1
    and 9 specifying the level of compression: 1 produces the least
    compression, and 9 produces the most compression. 3 is the default.
    """

    wbits = zlib.MAX_WBITS

    def __init__(self, filename, mode="rb", compresslevel=3):
        # This lock must be recursive, so that BufferedIOBase's
        # readline(), readlines() and writelines() don't deadlock.
        self._lock = RLock()
        self._fp = None
        self._closefp = False
        self._mode = _MODE_CLOSED
        self._pos = 0
        self._size = -1
        self.compresslevel = compresslevel

        if not isinstance(compresslevel, int) or not (1 <= compresslevel <= 9):
            raise ValueError(
                "'compresslevel' must be an integer "
                "between 1 and 9. You provided 'compresslevel={}'".format(compresslevel)
            )

        if mode == "rb":
            self._mode = _MODE_READ
            self._decompressor = zlib.decompressobj(self.wbits)
            self._buffer = b""
            self._buffer_offset = 0
        elif mode == "wb":
            self._mode = _MODE_WRITE
            self._compressor = zlib.compressobj(
                self.compresslevel, zlib.DEFLATED, self.wbits, zlib.DEF_MEM_LEVEL, 0
            )
        else:
            raise ValueError("Invalid mode: %r" % (mode,))

        if isinstance(filename, str):
            self._fp = io.open(filename, mode)
            self._closefp = True
        elif hasattr(filename, "read") or hasattr(filename, "write"):
            self._fp = filename
        else:
            raise TypeError("filename must be a str or bytes object, or a file")

    def close(self):
        """Flush and close the file.

        May be called more than once without error. Once the file is
        closed, any other operation on it will raise a ValueError.
        """
        with self._lock:
            if self._mode == _MODE_CLOSED:
                return
            try:
                if self._mode in (_MODE_READ, _MODE_READ_EOF):
                    self._decompressor = None
                elif self._mode == _MODE_WRITE:
                    self._fp.write(self._compressor.flush())
                    self._compressor = None
            finally:
                try:
                    if self._closefp:
                        self._fp.close()
                finally:
                    self._fp = None
                    self._closefp = False
                    self._mode = _MODE_CLOSED
                    self._buffer = b""
                    self._buffer_offset = 0

    @property
    def closed(self):
        """True if this file is closed."""
        return self._mode == _MODE_CLOSED

    def fileno(self):
        """Return the file descriptor for the underlying file."""
        self._check_not_closed()
        return self._fp.fileno()

    def seekable(self):
        """Return whether the file supports seeking."""
        return self.readable() and self._fp.seekable()

    def readable(self):
        """Return whether the file was opened for reading."""
        self._check_not_closed()
        return self._mode in (_MODE_READ, _MODE_READ_EOF)

    def writable(self):
        """Return whether the file was opened for writing."""
        self._check_not_closed()
        return self._mode == _MODE_WRITE

    # Mode-checking helper functions.

    def _check_not_closed(self):
        if self.closed:
            fname = getattr(self._fp, "name", None)
            msg = "I/O operation on closed file"
            if fname is not None:
                msg += " {}".format(fname)
            msg += "."
            raise ValueError(msg)

    def _check_can_read(self):
        if self._mode not in (_MODE_READ, _MODE_READ_EOF):
            self._check_not_closed()
            raise io.UnsupportedOperation("File not open for reading")

    def _check_can_write(self):
        if self._mode != _MODE_WRITE:
            self._check_not_closed()
            raise io.UnsupportedOperation("File not open for writing")

    def _check_can_seek(self):
        if self._mode not in (_MODE_READ, _MODE_READ_EOF):
            self._check_not_closed()
            raise io.UnsupportedOperation(
                "Seeking is only supported on files open for reading"
            )
        if not self._fp.seekable():
            raise io.UnsupportedOperation(
                "The underlying file object does not support seeking"
            )

    # Fill the readahead buffer if it is empty. Returns False on EOF.
    def _fill_buffer(self):
        if self._mode == _MODE_READ_EOF:
            return False
        # 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 self._buffer_offset == len(self._buffer):
            try:
                rawblock = self._decompressor.unused_data or self._fp.read(_BUFFER_SIZE)
                if not rawblock:
                    raise EOFError
            except EOFError:
                # End-of-stream marker and end of file. We're good.
                self._mode = _MODE_READ_EOF
                self._size = self._pos
                return False
            else:
                self._buffer = self._decompressor.decompress(rawblock)
            self._buffer_offset = 0
        return True

    # Read data until EOF.
    # If return_data is false, consume the data without returning it.
    def _read_all(self, return_data=True):
        # The loop assumes that _buffer_offset is 0. Ensure that this is true.
        self._buffer = self._buffer[self._buffer_offset :]
        self._buffer_offset = 0

        blocks = []
        while self._fill_buffer():
            if return_data:
                blocks.append(self._buffer)
            self._pos += len(self._buffer)
            self._buffer = b""
        if return_data:
            return b"".join(blocks)

    # Read a block of up to n bytes.
    # If return_data is false, consume the data without returning it.
    def _read_block(self, n_bytes, return_data=True):
        # If we have enough data buffered, return immediately.
        end = self._buffer_offset + n_bytes
        if end <= len(self._buffer):
            data = self._buffer[self._buffer_offset : end]
            self._buffer_offset = end
            self._pos += len(data)
            return data if return_data else None

        # The loop assumes that _buffer_offset is 0. Ensure that this is true.
        self._buffer = self._buffer[self._buffer_offset :]
        self._buffer_offset = 0

        blocks = []
        while n_bytes > 0 and self._fill_buffer():
            if n_bytes < len(self._buffer):
                data = self._buffer[:n_bytes]
                self._buffer_offset = n_bytes
            else:
                data = self._buffer
                self._buffer = b""
            if return_data:
                blocks.append(data)
            self._pos += len(data)
            n_bytes -= len(data)
        if return_data:
            return b"".join(blocks)

    def read(self, size=-1):
        """Read up to size uncompressed bytes from the file.

        If size is negative or omitted, read until EOF is reached.
        Returns b'' if the file is already at EOF.
        """
        with self._lock:
            self._check_can_read()
            if size == 0:
                return b""
            elif size < 0:
                return self._read_all()
            else:
                return self._read_block(size)

    def readinto(self, b):
        """Read up to len(b) bytes into b.

        Returns the number of bytes read (0 for EOF).
        """
        with self._lock:
            return io.BufferedIOBase.readinto(self, b)

    def write(self, data):
        """Write a byte string to the file.

        Returns the number of uncompressed bytes written, which is
        always len(data). Note that due to buffering, the file on disk
        may not reflect the data written until close() is called.
        """
        with self._lock:
            self._check_can_write()
            # Convert data type if called by io.BufferedWriter.
            if isinstance(data, memoryview):
                data = data.tobytes()

            compressed = self._compressor.compress(data)
            self._fp.write(compressed)
            self._pos += len(data)
            return len(data)

    # Rewind the file to the beginning of the data stream.
    def _rewind(self):
        self._fp.seek(0, 0)
        self._mode = _MODE_READ
        self._pos = 0
        self._decompressor = zlib.decompressobj(self.wbits)
        self._buffer = b""
        self._buffer_offset = 0

    def seek(self, offset, whence=0):
        """Change the file position.

        The new position is specified by offset, relative to the
        position indicated by whence. Values for whence are:

            0: start of stream (default); offset must not be negative
            1: current stream position
            2: end of stream; offset must not be positive

        Returns the new file position.

        Note that seeking is emulated, so depending on the parameters,
        this operation may be extremely slow.
        """
        with self._lock:
            self._check_can_seek()

            # Recalculate offset as an absolute file position.
            if whence == 0:
                pass
            elif whence == 1:
                offset = self._pos + offset
            elif whence == 2:
                # Seeking relative to EOF - we need to know the file's size.
                if self._size < 0:
                    self._read_all(return_data=False)
                offset = self._size + offset
            else:
                raise ValueError("Invalid value for whence: %s" % (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.
            self._read_block(offset, return_data=False)

            return self._pos

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


class ZlibCompressorWrapper(CompressorWrapper):
    def __init__(self):
        CompressorWrapper.__init__(
            self, obj=BinaryZlibFile, prefix=_ZLIB_PREFIX, extension=".z"
        )


class BinaryGzipFile(BinaryZlibFile):
    """A file object providing transparent gzip (de)compression.

    If filename is a str or bytes object, it gives the name
    of the file to be opened. Otherwise, it should be a file object,
    which will be used to read or write the compressed data.

    mode can be 'rb' for reading (default) or 'wb' for (over)writing

    If mode is 'wb', compresslevel can be a number between 1
    and 9 specifying the level of compression: 1 produces the least
    compression, and 9 produces the most compression. 3 is the default.
    """

    wbits = 31  # zlib compressor/decompressor wbits value for gzip format.


class GzipCompressorWrapper(CompressorWrapper):
    def __init__(self):
        CompressorWrapper.__init__(
            self, obj=BinaryGzipFile, prefix=_GZIP_PREFIX, extension=".gz"
        )