path: root/Lib/compression/zstd/_zstdfile.py

import io
from os import PathLike
from _zstd import (ZstdCompressor, ZstdDecompressor, ZstdError,
                   ZSTD_DStreamOutSize)
from compression._common import _streams

__all__ = ("ZstdFile", "open")

_MODE_CLOSED = 0
_MODE_READ = 1
_MODE_WRITE = 2


def _nbytes(dat, /):
    if isinstance(dat, (bytes, bytearray)):
        return len(dat)
    with memoryview(dat) as mv:
        return mv.nbytes


class ZstdFile(_streams.BaseStream):
    """A file-like object providing transparent Zstandard (de)compression.

    A ZstdFile can act as a wrapper for an existing file object, or refer
    directly to a named file on disk.

    ZstdFile provides a *binary* file interface. Data is read and returned as
    bytes, and may only be written to objects that support the Buffer Protocol.
    """

    FLUSH_BLOCK = ZstdCompressor.FLUSH_BLOCK
    FLUSH_FRAME = ZstdCompressor.FLUSH_FRAME

    def __init__(self, file, /, mode="r", *,
                 level=None, options=None, zstd_dict=None):
        """Open a Zstandard compressed file in binary mode.

        *file* can be either an file-like object, or a file name to open.

        *mode* can be "r" for reading (default), "w" for (over)writing, "x" for
        creating exclusively, or "a" for appending. These can equivalently be
        given as "rb", "wb", "xb" and "ab" respectively.

        *level* is an optional int specifying the compression level to use,
        or COMPRESSION_LEVEL_DEFAULT if not given.

        *options* is an optional dict for advanced compression parameters.
        See CompressionParameter and DecompressionParameter for the possible
        options.

        *zstd_dict* is an optional ZstdDict object, a pre-trained Zstandard
        dictionary. See train_dict() to train ZstdDict on sample data.
        """
        self._fp = None
        self._close_fp = False
        self._mode = _MODE_CLOSED
        self._buffer = None

        if not isinstance(mode, str):
            raise ValueError("mode must be a str")
        if options is not None and not isinstance(options, dict):
            raise TypeError("options must be a dict or None")
        mode = mode.removesuffix("b")  # handle rb, wb, xb, ab
        if mode == "r":
            if level is not None:
                raise TypeError("level is illegal in read mode")
            self._mode = _MODE_READ
        elif mode in {"w", "a", "x"}:
            if level is not None and not isinstance(level, int):
                raise TypeError("level must be int or None")
            self._mode = _MODE_WRITE
            self._compressor = ZstdCompressor(level=level, options=options,
                                              zstd_dict=zstd_dict)
            self._pos = 0
        else:
            raise ValueError(f"Invalid mode: {mode!r}")

        if isinstance(file, (str, bytes, PathLike)):
            self._fp = io.open(file, f'{mode}b')
            self._close_fp = True
        elif ((mode == 'r' and hasattr(file, "read"))
                or (mode != 'r' and hasattr(file, "write"))):
            self._fp = file
        else:
            raise TypeError("file must be a file-like object "
                            "or a str, bytes, or PathLike object")

        if self._mode == _MODE_READ:
            raw = _streams.DecompressReader(
                self._fp,
                ZstdDecompressor,
                zstd_dict=zstd_dict,
                options=options,
            )
            self._buffer = io.BufferedReader(raw)

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

        May be called multiple times. Once the file has been closed,
        any other operation on it will raise ValueError.
        """
        if self._fp is None:
            return
        try:
            if self._mode == _MODE_READ:
                if getattr(self, '_buffer', None):
                    self._buffer.close()
                    self._buffer = None
            elif self._mode == _MODE_WRITE:
                self.flush(self.FLUSH_FRAME)
                self._compressor = None
        finally:
            self._mode = _MODE_CLOSED
            try:
                if self._close_fp:
                    self._fp.close()
            finally:
                self._fp = None
                self._close_fp = False

    def write(self, data, /):
        """Write a bytes-like object *data* to the file.

        Returns the number of uncompressed bytes written, which is
        always the length of data in bytes. Note that due to buffering,
        the file on disk may not reflect the data written until .flush()
        or .close() is called.
        """
        self._check_can_write()

        length = _nbytes(data)

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

    def flush(self, mode=FLUSH_BLOCK):
        """Flush remaining data to the underlying stream.

        The mode argument can be FLUSH_BLOCK or FLUSH_FRAME. Abuse of this
        method will reduce compression ratio, use it only when necessary.

        If the program is interrupted afterwards, all data can be recovered.
        To ensure saving to disk, also need to use os.fsync(fd).

        This method does nothing in reading mode.
        """
        if self._mode == _MODE_READ:
            return
        self._check_not_closed()
        if mode not in {self.FLUSH_BLOCK, self.FLUSH_FRAME}:
            raise ValueError("Invalid mode argument, expected either "
                             "ZstdFile.FLUSH_FRAME or "
                             "ZstdFile.FLUSH_BLOCK")
        if self._compressor.last_mode == mode:
            return
        # Flush zstd block/frame, and write.
        data = self._compressor.flush(mode)
        self._fp.write(data)
        if hasattr(self._fp, "flush"):
            self._fp.flush()

    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.
        """
        if size is None:
            size = -1
        self._check_can_read()
        return self._buffer.read(size)

    def read1(self, size=-1):
        """Read up to size uncompressed bytes, while trying to avoid
        making multiple reads from the underlying stream. Reads up to a
        buffer's worth of data if size is negative.

        Returns b"" if the file is at EOF.
        """
        self._check_can_read()
        if size < 0:
            # Note this should *not* be io.DEFAULT_BUFFER_SIZE.
            # ZSTD_DStreamOutSize is the minimum amount to read guaranteeing
            # a full block is read.
            size = ZSTD_DStreamOutSize
        return self._buffer.read1(size)

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

        Returns the number of bytes read (0 for EOF).
        """
        self._check_can_read()
        return self._buffer.readinto(b)

    def readinto1(self, b):
        """Read bytes into b, while trying to avoid making multiple reads
        from the underlying stream.

        Returns the number of bytes read (0 for EOF).
        """
        self._check_can_read()
        return self._buffer.readinto1(b)

    def readline(self, size=-1):
        """Read a line of uncompressed bytes from the file.

        The terminating newline (if present) is retained. If size is
        non-negative, no more than size bytes will be read (in which
        case the line may be incomplete). Returns b'' if already at EOF.
        """
        self._check_can_read()
        return self._buffer.readline(size)

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

        The new position is specified by offset, relative to the
        position indicated by whence. Possible 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 arguments,
        this operation may be extremely slow.
        """
        self._check_can_read()

        # BufferedReader.seek() checks seekable
        return self._buffer.seek(offset, whence)

    def peek(self, size=-1):
        """Return buffered data without advancing the file position.

        Always returns at least one byte of data, unless at EOF.
        The exact number of bytes returned is unspecified.
        """
        # Relies on the undocumented fact that BufferedReader.peek() always
        # returns at least one byte (except at EOF)
        self._check_can_read()
        return self._buffer.peek(size)

    def __next__(self):
        if ret := self._buffer.readline():
            return ret
        raise StopIteration

    def tell(self):
        """Return the current file position."""
        self._check_not_closed()
        if self._mode == _MODE_READ:
            return self._buffer.tell()
        elif self._mode == _MODE_WRITE:
            return self._pos

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

    @property
    def name(self):
        self._check_not_closed()
        return self._fp.name

    @property
    def mode(self):
        return 'wb' if self._mode == _MODE_WRITE else 'rb'

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

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

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

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


def open(file, /, mode="rb", *, level=None, options=None, zstd_dict=None,
         encoding=None, errors=None, newline=None):
    """Open a Zstandard compressed file in binary or text mode.

    file can be either a file name (given as a str, bytes, or PathLike object),
    in which case the named file is opened, or it can be an existing file object
    to read from or write to.

    The mode parameter can be "r", "rb" (default), "w", "wb", "x", "xb", "a",
    "ab" for binary mode, or "rt", "wt", "xt", "at" for text mode.

    The level, options, and zstd_dict parameters specify the settings the same
    as ZstdFile.

    When using read mode (decompression), the options parameter is a dict
    representing advanced decompression options. The level parameter is not
    supported in this case. When using write mode (compression), only one of
    level, an int representing the compression level, or options, a dict
    representing advanced compression options, may be passed. In both modes,
    zstd_dict is a ZstdDict instance containing a trained Zstandard dictionary.

    For binary mode, this function is equivalent to the ZstdFile constructor:
    ZstdFile(filename, mode, ...). In this case, the encoding, errors and
    newline parameters must not be provided.

    For text mode, an ZstdFile object is created, and wrapped in an
    io.TextIOWrapper instance with the specified encoding, error handling
    behavior, and line ending(s).
    """

    text_mode = "t" in mode
    mode = mode.replace("t", "")

    if text_mode:
        if "b" in mode:
            raise ValueError(f"Invalid mode: {mode!r}")
    else:
        if encoding is not None:
            raise ValueError("Argument 'encoding' not supported in binary mode")
        if errors is not None:
            raise ValueError("Argument 'errors' not supported in binary mode")
        if newline is not None:
            raise ValueError("Argument 'newline' not supported in binary mode")

    binary_file = ZstdFile(file, mode, level=level, options=options,
                           zstd_dict=zstd_dict)

    if text_mode:
        return io.TextIOWrapper(binary_file, encoding, errors, newline)
    else:
        return binary_file