| import io |
| from os import PathLike |
| from _zstd import ZstdCompressor, ZstdDecompressor, 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 |
