from array import array from bisect import bisect_right from os.path import isfile from struct import Struct from warnings import warn from pyzstd.zstdfile import ZstdDecompressReader, ZstdFile, \ _MODE_CLOSED, _MODE_READ, _MODE_WRITE, \ PathLike, io __all__ = ('SeekableFormatError', 'SeekableZstdFile') class SeekableFormatError(Exception): 'An error related to Zstandard Seekable Format.' def __init__(self, msg): super().__init__('Zstandard Seekable Format error: ' + msg) __doc__ = '''\ Zstandard Seekable Format (Ver 0.1.0, Apr 2017) Square brackets are used to indicate optional fields. All numeric fields are little-endian unless specified otherwise. A. Seek table is a skippable frame at the end of file: Magic_Number Frame_Size [Seek_Table_Entries] Seek_Table_Footer 4 bytes 4 bytes 8-12 bytes each 9 bytes Magic_Number must be 0x184D2A5E. B. Seek_Table_Entries: Compressed_Size Decompressed_Size [Checksum] 4 bytes 4 bytes 4 bytes Checksum is optional. C. Seek_Table_Footer: Number_Of_Frames Seek_Table_Descriptor Seekable_Magic_Number 4 bytes 1 byte 4 bytes Seekable_Magic_Number must be 0x8F92EAB1. D. Seek_Table_Descriptor: Bit_number Field_name 7 Checksum_Flag 6-2 Reserved_Bits (should ensure they are set to 0) 1-0 Unused_Bits (should not interpret these bits)''' __format_version__ = '0.1.0' class SeekTable: _s_2uint32 = Struct(' fsize - skippable_frame_size: msg = ('Wrong seek table. Since index %d frame (0-based), ' 'the cumulated compressed size is greater than ' 'file size.') % idx raise SeekableFormatError(msg) # Check format if self._full_c_size != fsize - skippable_frame_size: raise SeekableFormatError('The cumulated compressed size is wrong') # Parsed successfully, save for future use. self._seek_frame_size = skippable_frame_size self._file_size = fsize # Find frame index by decompressed position def index_by_dpos(self, pos): # Array's first item is 0, so need this. if pos < 0: pos = 0 i = bisect_right(self._cumulated_d_size, pos) if i != self._frames_count + 1: return i else: # None means >= EOF return None def get_frame_sizes(self, i): return (self._cumulated_c_size[i-1], self._cumulated_d_size[i-1]) def get_full_c_size(self): return self._full_c_size def get_full_d_size(self): return self._full_d_size # Merge the seek table to max_frames frames. # The format allows up to 0xFFFF_FFFF frames. When frames # number exceeds, use this method to merge. def _merge_frames(self, max_frames): if self._frames_count <= max_frames: return # Clear the table arr = self._frames a, b = divmod(self._frames_count, max_frames) self._clear_seek_table() # Merge frames pos = 0 for i in range(max_frames): # Slice length length = (a + (1 if i < b else 0)) * 2 # Merge c_size = 0 d_size = 0 for j in range(pos, pos+length, 2): c_size += arr[j] d_size += arr[j+1] self.append_entry(c_size, d_size) pos += length def write_seek_table(self, fp): # Exceeded format limit if self._frames_count > 0xFFFFFFFF: # Emit a warning warn(('SeekableZstdFile\'s seek table has %d entries, ' 'which exceeds the maximal value allowed by ' 'Zstandard Seekable Format (0xFFFFFFFF). The ' 'entries will be merged into 0xFFFFFFFF entries, ' 'this may reduce seeking performance.') % self._frames_count, RuntimeWarning, 3) # Merge frames self._merge_frames(0xFFFFFFFF) # The skippable frame offset = 0 size = 17 + 8 * self._frames_count ba = bytearray(size) # Header self._s_2uint32.pack_into(ba, offset, 0x184D2A5E, size-8) offset += 8 # Entries for i in range(0, len(self._frames), 2): self._s_2uint32.pack_into(ba, offset, self._frames[i], self._frames[i+1]) offset += 8 # Footer self._s_footer.pack_into(ba, offset, self._frames_count, 0, 0x8F92EAB1) # Write fp.write(ba) @property def seek_frame_size(self): return self._seek_frame_size @property def file_size(self): return self._file_size def __len__(self): return self._frames_count def get_info(self): return (self._frames_count, self._full_c_size, self._full_d_size) class SeekableDecompressReader(ZstdDecompressReader): def __init__(self, fp, zstd_dict, option, read_size): # Check fp readable/seekable if not hasattr(fp, 'readable') or not hasattr(fp, "seekable"): raise TypeError( ("In SeekableZstdFile's reading mode, the file object should " "have .readable()/.seekable() methods.")) if not fp.readable(): raise TypeError( ("In SeekableZstdFile's reading mode, the file object should " "be readable.")) if not fp.seekable(): raise TypeError( ("In SeekableZstdFile's reading mode, the file object should " "be seekable. If the file object is not seekable, it can be " "read sequentially using ZstdFile class.")) # Load seek table self._seek_table = SeekTable(read_mode=True) self._seek_table.load_seek_table(fp, seek_to_0=True) # Initialize super() super().__init__(fp, zstd_dict, option, read_size) self._decomp.size = self._seek_table.get_full_d_size() # super().seekable() returns self._fp.seekable(). # Seekable has been checked in .__init__() method. # BufferedReader.seek() checks this in each invoke, if self._fp.seekable() # becomes False at runtime, .seek() method just raise OSError instead of # io.UnsupportedOperation. def seekable(self): return True # If the new position is within BufferedReader's buffer, # this method may not be called. def seek(self, offset, whence=0): # offset is absolute file position if whence == 0: # SEEK_SET pass elif whence == 1: # SEEK_CUR offset = self._decomp.pos + offset elif whence == 2: # SEEK_END offset = self._decomp.size + offset else: raise ValueError("Invalid value for whence: {}".format(whence)) # Get new frame index new_frame = self._seek_table.index_by_dpos(offset) # offset >= EOF if new_frame is None: self._decomp.eof = True self._decomp.pos = self._decomp.size self._fp.seek(self._seek_table.file_size) return self._decomp.pos # Prepare to jump old_frame = self._seek_table.index_by_dpos(self._decomp.pos) c_pos, d_pos = self._seek_table.get_frame_sizes(new_frame) # If at P1, seeking to P2 will unnecessarily read the skippable # frame. So check self._fp position to skip the skippable frame. # |--data1--|--skippable--|--data2--| # cpos: ^P1 # dpos: ^P1 ^P2 if new_frame == old_frame and \ offset >= self._decomp.pos and \ self._fp.tell() >= c_pos: pass else: # Jump self._decomp.eof = False self._decomp.pos = d_pos self._decomp.reset_session() self._fp.seek(c_pos) # offset is bytes number to skip forward offset -= self._decomp.pos # If offset <= 0, .forward() method does nothing. self._decomp.forward(offset) return self._decomp.pos def get_seek_table_info(self): return self._seek_table.get_info() # Compared to ZstdFile class, it's important to handle the seekable # of underlying file object carefully. Need to check seekable in # each situation. For example, there may be a CD-R file system that # is seekable when reading, but not seekable when appending. class SeekableZstdFile(ZstdFile): """This class can only create/write/read Zstandard Seekable Format file, or read 0-size file. It provides relatively fast seeking ability in read mode. """ # The format uses uint32_t for compressed/decompressed sizes. If flush # block a lot, compressed_size may exceed the limit, so set a max size. FRAME_MAX_C_SIZE = 2*1024*1024*1024 # Zstd seekable format's example code also use 1GiB as max content size. FRAME_MAX_D_SIZE = 1*1024*1024*1024 _READER_CLASS = SeekableDecompressReader def __init__(self, filename, mode="r", *, level_or_option=None, zstd_dict=None, read_size=131075, write_size=131591, max_frame_content_size=1024*1024*1024): """Open a Zstandard Seekable Format file in binary mode. In read mode, the file can be 0-size file. filename can be either an actual 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. 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. In append mode ("a" or "ab"), filename argument can't be a file object, please use file path. Parameters level_or_option: When it's an int object, it represents compression level. When it's a dict object, it contains advanced compression parameters. Note, in read mode (decompression), it can only be a dict object, that represents decompression option. It doesn't support int type compression level in this case. zstd_dict: A ZstdDict object, pre-trained dictionary for compression / decompression. read_size: In reading mode, this is bytes number that read from the underlying file object each time, default value is zstd's recommended value. If use with Network File System, increasing it may get better performance. write_size: In writing modes, this is output buffer's size, default value is zstd's recommended value. If use with Network File System, increasing it may get better performance. max_frame_content_size: In write/append modes (compression), when the uncompressed data size reaches max_frame_content_size, a frame is generated automatically. If the size is small, it will increase seeking speed, but reduce compression ratio. If the size is large, it will reduce seeking speed, but increase compression ratio. You can also manually generate a frame using f.flush(f.FLUSH_FRAME). """ # For self.close() self._write_in_close = False # For super().close() self._fp = None self._closefp = False self._mode = _MODE_CLOSED if mode in ("r", "rb"): # Specified max_frame_content_size argument if max_frame_content_size != 1024*1024*1024: raise ValueError(('max_frame_content_size argument is only ' 'valid in write modes (compression).')) elif mode in ("w", "wb", "a", "ab", "x", "xb"): if not (0 < max_frame_content_size <= self.FRAME_MAX_D_SIZE): raise ValueError( ('max_frame_content_size argument should be ' '0 < value <= %d, provided value is %d.') % \ (self.FRAME_MAX_D_SIZE, max_frame_content_size)) # For seekable format self._max_frame_content_size = max_frame_content_size self._reset_frame_sizes() self._seek_table = SeekTable(read_mode=False) # Load seek table in append mode if mode in ("a", "ab"): if not isinstance(filename, (str, bytes, PathLike)): raise TypeError( ("In append mode ('a', 'ab'), " "SeekableZstdFile.__init__() method can't " "accept file object as filename argument. " "Please use file path (str/bytes/PathLike).")) # Load seek table if file exists if isfile(filename): with io.open(filename, "rb") as f: if not hasattr(f, "seekable") or not f.seekable(): raise TypeError( ("In SeekableZstdFile's append mode " "('a', 'ab'), the opened 'rb' file " "object should be seekable.")) self._seek_table.load_seek_table(f, seek_to_0=False) super().__init__(filename, mode, level_or_option=level_or_option, zstd_dict=zstd_dict, read_size=read_size, write_size=write_size) # Overwrite seek table in append mode if mode in ("a", "ab"): if self._fp.seekable(): self._fp.seek(self._seek_table.get_full_c_size()) # Necessary if the current table has many (0, 0) entries self._fp.truncate() else: # Add the seek table frame self._seek_table.append_entry( self._seek_table.seek_frame_size, 0) # Emit a warning warn(("SeekableZstdFile is opened in append mode " "('a', 'ab'), but the underlying file object " "is not seekable. Therefore the seek table (a " "zstd skippable frame) at the end of the file " "can't be overwritten. Each time open such file " "in append mode, it will waste some storage " "space. %d bytes were wasted this time.") % \ self._seek_table.seek_frame_size, RuntimeWarning, 2) # Initialized successfully self._write_in_close = (self._mode == _MODE_WRITE) def _reset_frame_sizes(self): self._current_c_size = 0 self._current_d_size = 0 self._left_d_size = self._max_frame_content_size 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. """ try: if self._write_in_close: try: self.flush(self.FLUSH_FRAME) self._seek_table.write_seek_table(self._fp) finally: # For multiple calls to .close() self._write_in_close = False finally: # Clear write mode's seek table. # Put here for failures in/after super().__init__(). self._seek_table = None super().close() def write(self, data): """Write a bytes-like object 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. """ if self._mode != _MODE_WRITE: self._check_mode(_MODE_WRITE) # Accept any data that supports the buffer protocol. # And memoryview's subview is faster than slice. with memoryview(data) as view, view.cast('B') as byte_view: nbytes = byte_view.nbytes pos = 0 while nbytes > 0: # Write size write_size = min(nbytes, self._left_d_size) # Use inserted super().write() method, to prevent # self._fp.tell() from reporting incorrect position. # ------------------------- # super().write() begin # ------------------------- # Compress & write _, output_size = self._writer.write( byte_view[pos:pos+write_size]) self._pos += write_size # ----------------------- # super().write() end # ----------------------- pos += write_size nbytes -= write_size # Cumulate self._current_c_size += output_size self._current_d_size += write_size self._left_d_size -= write_size # Should flush a frame if self._left_d_size == 0 or \ self._current_c_size >= self.FRAME_MAX_C_SIZE: self.flush(self.FLUSH_FRAME) return pos def flush(self, mode=ZstdFile.FLUSH_BLOCK): """Flush remaining data to the underlying stream. The mode argument can be ZstdFile.FLUSH_BLOCK, ZstdFile.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_WRITE: # Like IOBase.flush(), do nothing in reading mode. # TextIOWrapper.close() relies on this behavior. if self._mode == _MODE_READ: return # Closed, raise ValueError. self._check_mode() # Use inserted super().flush() method, to prevent # self._fp.tell() from reporting incorrect position. # ------------------------- # super().flush() begin # ------------------------- # Flush zstd block/frame, and write. _, output_size = self._writer.flush(mode) # ----------------------- # super().flush() end # ----------------------- # Cumulate self._current_c_size += output_size # self._current_d_size += 0 # self._left_d_size -= 0 if mode == self.FLUSH_FRAME and \ self._current_c_size != 0: # Add an entry to seek table self._seek_table.append_entry(self._current_c_size, self._current_d_size) self._reset_frame_sizes() @property def seek_table_info(self): """A tuple: (frames_number, compressed_size, decompressed_size) 1, Frames_number and compressed_size don't count the seek table frame (a zstd skippable frame at the end of the file). 2, In write modes, the part of data that has not been flushed to frames is not counted. 3, If the SeekableZstdFile object is closed, it's None. """ if self._mode == _MODE_WRITE: return self._seek_table.get_info() elif self._mode == _MODE_READ: return self._buffer.raw.get_seek_table_info() else: # Closed return None @staticmethod def is_seekable_format_file(filename): """Check if a file is Zstandard Seekable Format file or 0-size file. It parses the seek table at the end of the file, returns True if no format error. filename can be either a file path (str/bytes/PathLike), or can be an existing file object in reading mode. """ # Check argument if isinstance(filename, (str, bytes, PathLike)): fp = io.open(filename, 'rb') is_file_path = True elif hasattr(filename, 'readable') and filename.readable() and \ hasattr(filename, "seekable") and filename.seekable(): fp = filename is_file_path = False orig_pos = fp.tell() else: raise TypeError( ('filename argument should be a str/bytes/PathLike object, ' 'or a file object that is readable and seekable.')) # Write mode uses less RAM table = SeekTable(read_mode=False) try: # Read/Parse the seek table table.load_seek_table(fp, seek_to_0=False) except SeekableFormatError: ret = False else: ret = True finally: if is_file_path: fp.close() else: fp.seek(orig_pos) return ret