def _ensure_dstream(self):
if self._dstream:
zresult = lib.ZSTD_resetDStream(self._dstream)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not reset DStream: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
return
self._dstream = lib.ZSTD_createDStream()
if self._dstream == ffi.NULL:
raise MemoryError()
self._dstream = ffi.gc(self._dstream, lib.ZSTD_freeDStream)
if self._dict_data:
zresult = lib.ZSTD_initDStream_usingDict(
self._dstream, self._dict_data.as_bytes(),
len(self._dict_data))
else:
zresult = lib.ZSTD_initDStream(self._dstream)
if lib.ZSTD_isError(zresult):
self._dstream = None
raise ZstdError('could not initialize DStream: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
def __init__(self, dict_data=None):
self._dict_data = dict_data
dctx = lib.ZSTD_createDCtx()
if dctx == ffi.NULL:
raise MemoryError()
self._refdctx = ffi.gc(dctx, lib.ZSTD_freeDCtx)
self._dstream = None
def _get_cstream(self):
cstream = lib.ZSTD_createCStream()
cstream = ffi.gc(cstream, lib.ZSTD_freeCStream)
res = lib.ZSTD_initCStream(cstream, self._compression_level)
if lib.ZSTD_isError(res):
raise Exception('cannot init CStream: %s' %
lib.ZSTD_getErrorName(res))
return cstream
def __init__(self,
level=3,
dict_data=None,
compression_params=None,
write_checksum=False,
write_content_size=False,
write_dict_id=True,
threads=0):
if level < 1:
raise ValueError('level must be greater than 0')
elif level > lib.ZSTD_maxCLevel():
raise ValueError('level must be less than %d' %
lib.ZSTD_maxCLevel())
if threads < 0:
threads = _cpu_count()
self._compression_level = level
self._dict_data = dict_data
self._cparams = compression_params
self._fparams = ffi.new('ZSTD_frameParameters *')[0]
self._fparams.checksumFlag = write_checksum
self._fparams.contentSizeFlag = write_content_size
self._fparams.noDictIDFlag = not write_dict_id
if threads:
cctx = lib.ZSTDMT_createCCtx(threads)
if cctx == ffi.NULL:
raise MemoryError()
self._cctx = ffi.gc(cctx, lib.ZSTDMT_freeCCtx)
self._multithreaded = True
else:
cctx = lib.ZSTD_createCCtx()
if cctx == ffi.NULL:
raise MemoryError()
self._cctx = ffi.gc(cctx, lib.ZSTD_freeCCtx)
self._multithreaded = False
self._cstream = None
def _ensure_cstream(self, size):
if self._cstream:
zresult = lib.ZSTD_resetCStream(self._cstream, size)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not reset CStream: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
return
cstream = lib.ZSTD_createCStream()
if cstream == ffi.NULL:
raise MemoryError()
cstream = ffi.gc(cstream, lib.ZSTD_freeCStream)
dict_data = ffi.NULL
dict_size = 0
if self._dict_data:
dict_data = self._dict_data.as_bytes()
dict_size = len(self._dict_data)
zparams = ffi.new('ZSTD_parameters *')[0]
if self._cparams:
zparams.cParams = self._cparams.as_compression_parameters()
else:
zparams.cParams = lib.ZSTD_getCParams(self._compression_level,
size, dict_size)
zparams.fParams = self._fparams
zresult = lib.ZSTD_initCStream_advanced(cstream, dict_data, dict_size,
zparams, size)
if lib.ZSTD_isError(zresult):
raise Exception('cannot init CStream: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
self._cstream = cstream
def decompress_content_dict_chain(self, frames):
if not isinstance(frames, list):
raise TypeError('argument must be a list')
if not frames:
raise ValueError('empty input chain')
# First chunk should not be using a dictionary. We handle it specially.
chunk = frames[0]
if not isinstance(chunk, bytes_type):
raise ValueError('chunk 0 must be bytes')
# All chunks should be zstd frames and should have content size set.
chunk_buffer = ffi.from_buffer(chunk)
params = ffi.new('ZSTD_frameParams *')
zresult = lib.ZSTD_getFrameParams(params, chunk_buffer,
len(chunk_buffer))
if lib.ZSTD_isError(zresult):
raise ValueError('chunk 0 is not a valid zstd frame')
elif zresult:
raise ValueError('chunk 0 is too small to contain a zstd frame')
if not params.frameContentSize:
raise ValueError('chunk 0 missing content size in frame')
dctx = lib.ZSTD_createDCtx()
if dctx == ffi.NULL:
raise MemoryError()
dctx = ffi.gc(dctx, lib.ZSTD_freeDCtx)
last_buffer = ffi.new('char[]', params.frameContentSize)
zresult = lib.ZSTD_decompressDCtx(dctx, last_buffer, len(last_buffer),
chunk_buffer, len(chunk_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('could not decompress chunk 0: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
# Special case of chain length of 1
if len(frames) == 1:
return ffi.buffer(last_buffer, len(last_buffer))[:]
i = 1
while i < len(frames):
chunk = frames[i]
if not isinstance(chunk, bytes_type):
raise ValueError('chunk %d must be bytes' % i)
chunk_buffer = ffi.from_buffer(chunk)
zresult = lib.ZSTD_getFrameParams(params, chunk_buffer,
len(chunk_buffer))
if lib.ZSTD_isError(zresult):
raise ValueError('chunk %d is not a valid zstd frame' % i)
elif zresult:
raise ValueError(
'chunk %d is too small to contain a zstd frame' % i)
if not params.frameContentSize:
raise ValueError('chunk %d missing content size in frame' % i)
dest_buffer = ffi.new('char[]', params.frameContentSize)
zresult = lib.ZSTD_decompress_usingDict(dctx, dest_buffer,
len(dest_buffer),
chunk_buffer,
len(chunk_buffer),
last_buffer,
len(last_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('could not decompress chunk %d' % i)
last_buffer = dest_buffer
i += 1
return ffi.buffer(last_buffer, len(last_buffer))[:]