grep: restore pre-9ef10437bb88 behavior, enable wdir search by tweakdefaults
Unfortunately, python-hglib relies on the original grep behavior and is
documented as such. Even though we agreed to introduce the BC, we shouldn't
break existing libraries.
So this patch flips the default again and move the new default to
ui.tweakdefaults. We could instead use HGPLAIN to turn this flag off, but
that would be rather confusing as the old/new behaviors are quite different.
Differential Revision: https://phab.mercurial-scm.org/D3919
# Copyright (c) 2016-present, Gregory Szorc
# All rights reserved.
#
# This software may be modified and distributed under the terms
# of the BSD license. See the LICENSE file for details.
"""Python interface to the Zstandard (zstd) compression library."""
from __future__ import absolute_import, unicode_literals
# This should match what the C extension exports.
__all__ = [
#'BufferSegment',
#'BufferSegments',
#'BufferWithSegments',
#'BufferWithSegmentsCollection',
'CompressionParameters',
'ZstdCompressionDict',
'ZstdCompressionParameters',
'ZstdCompressor',
'ZstdError',
'ZstdDecompressor',
'FrameParameters',
'estimate_decompression_context_size',
'frame_content_size',
'frame_header_size',
'get_frame_parameters',
'train_dictionary',
# Constants.
'COMPRESSOBJ_FLUSH_FINISH',
'COMPRESSOBJ_FLUSH_BLOCK',
'ZSTD_VERSION',
'FRAME_HEADER',
'CONTENTSIZE_UNKNOWN',
'CONTENTSIZE_ERROR',
'MAX_COMPRESSION_LEVEL',
'COMPRESSION_RECOMMENDED_INPUT_SIZE',
'COMPRESSION_RECOMMENDED_OUTPUT_SIZE',
'DECOMPRESSION_RECOMMENDED_INPUT_SIZE',
'DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE',
'MAGIC_NUMBER',
'WINDOWLOG_MIN',
'WINDOWLOG_MAX',
'CHAINLOG_MIN',
'CHAINLOG_MAX',
'HASHLOG_MIN',
'HASHLOG_MAX',
'HASHLOG3_MAX',
'SEARCHLOG_MIN',
'SEARCHLOG_MAX',
'SEARCHLENGTH_MIN',
'SEARCHLENGTH_MAX',
'TARGETLENGTH_MIN',
'LDM_MINMATCH_MIN',
'LDM_MINMATCH_MAX',
'LDM_BUCKETSIZELOG_MAX',
'STRATEGY_FAST',
'STRATEGY_DFAST',
'STRATEGY_GREEDY',
'STRATEGY_LAZY',
'STRATEGY_LAZY2',
'STRATEGY_BTLAZY2',
'STRATEGY_BTOPT',
'STRATEGY_BTULTRA',
'DICT_TYPE_AUTO',
'DICT_TYPE_RAWCONTENT',
'DICT_TYPE_FULLDICT',
'FORMAT_ZSTD1',
'FORMAT_ZSTD1_MAGICLESS',
]
import io
import os
import sys
from _zstd_cffi import (
ffi,
lib,
)
if sys.version_info[0] == 2:
bytes_type = str
int_type = long
else:
bytes_type = bytes
int_type = int
COMPRESSION_RECOMMENDED_INPUT_SIZE = lib.ZSTD_CStreamInSize()
COMPRESSION_RECOMMENDED_OUTPUT_SIZE = lib.ZSTD_CStreamOutSize()
DECOMPRESSION_RECOMMENDED_INPUT_SIZE = lib.ZSTD_DStreamInSize()
DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE = lib.ZSTD_DStreamOutSize()
new_nonzero = ffi.new_allocator(should_clear_after_alloc=False)
MAX_COMPRESSION_LEVEL = lib.ZSTD_maxCLevel()
MAGIC_NUMBER = lib.ZSTD_MAGICNUMBER
FRAME_HEADER = b'\x28\xb5\x2f\xfd'
CONTENTSIZE_UNKNOWN = lib.ZSTD_CONTENTSIZE_UNKNOWN
CONTENTSIZE_ERROR = lib.ZSTD_CONTENTSIZE_ERROR
ZSTD_VERSION = (lib.ZSTD_VERSION_MAJOR, lib.ZSTD_VERSION_MINOR, lib.ZSTD_VERSION_RELEASE)
WINDOWLOG_MIN = lib.ZSTD_WINDOWLOG_MIN
WINDOWLOG_MAX = lib.ZSTD_WINDOWLOG_MAX
CHAINLOG_MIN = lib.ZSTD_CHAINLOG_MIN
CHAINLOG_MAX = lib.ZSTD_CHAINLOG_MAX
HASHLOG_MIN = lib.ZSTD_HASHLOG_MIN
HASHLOG_MAX = lib.ZSTD_HASHLOG_MAX
HASHLOG3_MAX = lib.ZSTD_HASHLOG3_MAX
SEARCHLOG_MIN = lib.ZSTD_SEARCHLOG_MIN
SEARCHLOG_MAX = lib.ZSTD_SEARCHLOG_MAX
SEARCHLENGTH_MIN = lib.ZSTD_SEARCHLENGTH_MIN
SEARCHLENGTH_MAX = lib.ZSTD_SEARCHLENGTH_MAX
TARGETLENGTH_MIN = lib.ZSTD_TARGETLENGTH_MIN
LDM_MINMATCH_MIN = lib.ZSTD_LDM_MINMATCH_MIN
LDM_MINMATCH_MAX = lib.ZSTD_LDM_MINMATCH_MAX
LDM_BUCKETSIZELOG_MAX = lib.ZSTD_LDM_BUCKETSIZELOG_MAX
STRATEGY_FAST = lib.ZSTD_fast
STRATEGY_DFAST = lib.ZSTD_dfast
STRATEGY_GREEDY = lib.ZSTD_greedy
STRATEGY_LAZY = lib.ZSTD_lazy
STRATEGY_LAZY2 = lib.ZSTD_lazy2
STRATEGY_BTLAZY2 = lib.ZSTD_btlazy2
STRATEGY_BTOPT = lib.ZSTD_btopt
STRATEGY_BTULTRA = lib.ZSTD_btultra
DICT_TYPE_AUTO = lib.ZSTD_dct_auto
DICT_TYPE_RAWCONTENT = lib.ZSTD_dct_rawContent
DICT_TYPE_FULLDICT = lib.ZSTD_dct_fullDict
FORMAT_ZSTD1 = lib.ZSTD_f_zstd1
FORMAT_ZSTD1_MAGICLESS = lib.ZSTD_f_zstd1_magicless
COMPRESSOBJ_FLUSH_FINISH = 0
COMPRESSOBJ_FLUSH_BLOCK = 1
def _cpu_count():
# os.cpu_count() was introducd in Python 3.4.
try:
return os.cpu_count() or 0
except AttributeError:
pass
# Linux.
try:
if sys.version_info[0] == 2:
return os.sysconf(b'SC_NPROCESSORS_ONLN')
else:
return os.sysconf(u'SC_NPROCESSORS_ONLN')
except (AttributeError, ValueError):
pass
# TODO implement on other platforms.
return 0
class ZstdError(Exception):
pass
def _zstd_error(zresult):
# Resolves to bytes on Python 2 and 3. We use the string for formatting
# into error messages, which will be literal unicode. So convert it to
# unicode.
return ffi.string(lib.ZSTD_getErrorName(zresult)).decode('utf-8')
def _make_cctx_params(params):
res = lib.ZSTD_createCCtxParams()
if res == ffi.NULL:
raise MemoryError()
res = ffi.gc(res, lib.ZSTD_freeCCtxParams)
attrs = [
(lib.ZSTD_p_format, params.format),
(lib.ZSTD_p_compressionLevel, params.compression_level),
(lib.ZSTD_p_windowLog, params.window_log),
(lib.ZSTD_p_hashLog, params.hash_log),
(lib.ZSTD_p_chainLog, params.chain_log),
(lib.ZSTD_p_searchLog, params.search_log),
(lib.ZSTD_p_minMatch, params.min_match),
(lib.ZSTD_p_targetLength, params.target_length),
(lib.ZSTD_p_compressionStrategy, params.compression_strategy),
(lib.ZSTD_p_contentSizeFlag, params.write_content_size),
(lib.ZSTD_p_checksumFlag, params.write_checksum),
(lib.ZSTD_p_dictIDFlag, params.write_dict_id),
(lib.ZSTD_p_nbWorkers, params.threads),
(lib.ZSTD_p_jobSize, params.job_size),
(lib.ZSTD_p_overlapSizeLog, params.overlap_size_log),
(lib.ZSTD_p_compressLiterals, params.compress_literals),
(lib.ZSTD_p_forceMaxWindow, params.force_max_window),
(lib.ZSTD_p_enableLongDistanceMatching, params.enable_ldm),
(lib.ZSTD_p_ldmHashLog, params.ldm_hash_log),
(lib.ZSTD_p_ldmMinMatch, params.ldm_min_match),
(lib.ZSTD_p_ldmBucketSizeLog, params.ldm_bucket_size_log),
(lib.ZSTD_p_ldmHashEveryLog, params.ldm_hash_every_log),
]
for param, value in attrs:
_set_compression_parameter(res, param, value)
return res
class ZstdCompressionParameters(object):
@staticmethod
def from_level(level, source_size=0, dict_size=0, **kwargs):
params = lib.ZSTD_getCParams(level, source_size, dict_size)
args = {
'window_log': 'windowLog',
'chain_log': 'chainLog',
'hash_log': 'hashLog',
'search_log': 'searchLog',
'min_match': 'searchLength',
'target_length': 'targetLength',
'compression_strategy': 'strategy',
}
for arg, attr in args.items():
if arg not in kwargs:
kwargs[arg] = getattr(params, attr)
if 'compress_literals' not in kwargs:
kwargs['compress_literals'] = 1 if level >= 0 else 0
return ZstdCompressionParameters(**kwargs)
def __init__(self, format=0, compression_level=0, window_log=0, hash_log=0,
chain_log=0, search_log=0, min_match=0, target_length=0,
compression_strategy=0, write_content_size=1, write_checksum=0,
write_dict_id=0, job_size=0, overlap_size_log=0,
force_max_window=0, enable_ldm=0, ldm_hash_log=0,
ldm_min_match=0, ldm_bucket_size_log=0, ldm_hash_every_log=0,
threads=0, compress_literals=None):
if threads < 0:
threads = _cpu_count()
if compress_literals is None:
compress_literals = compression_level >= 0
self.format = format
self.compression_level = compression_level
self.window_log = window_log
self.hash_log = hash_log
self.chain_log = chain_log
self.search_log = search_log
self.min_match = min_match
self.target_length = target_length
self.compression_strategy = compression_strategy
self.write_content_size = write_content_size
self.write_checksum = write_checksum
self.write_dict_id = write_dict_id
self.job_size = job_size
self.overlap_size_log = overlap_size_log
self.compress_literals = compress_literals
self.force_max_window = force_max_window
self.enable_ldm = enable_ldm
self.ldm_hash_log = ldm_hash_log
self.ldm_min_match = ldm_min_match
self.ldm_bucket_size_log = ldm_bucket_size_log
self.ldm_hash_every_log = ldm_hash_every_log
self.threads = threads
self.params = _make_cctx_params(self)
def estimated_compression_context_size(self):
return lib.ZSTD_estimateCCtxSize_usingCCtxParams(self.params)
CompressionParameters = ZstdCompressionParameters
def estimate_decompression_context_size():
return lib.ZSTD_estimateDCtxSize()
def _set_compression_parameter(params, param, value):
zresult = lib.ZSTD_CCtxParam_setParameter(params, param,
ffi.cast('unsigned', value))
if lib.ZSTD_isError(zresult):
raise ZstdError('unable to set compression context parameter: %s' %
_zstd_error(zresult))
class ZstdCompressionWriter(object):
def __init__(self, compressor, writer, source_size, write_size):
self._compressor = compressor
self._writer = writer
self._source_size = source_size
self._write_size = write_size
self._entered = False
self._bytes_compressed = 0
def __enter__(self):
if self._entered:
raise ZstdError('cannot __enter__ multiple times')
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._compressor._cctx,
self._source_size)
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
self._entered = True
return self
def __exit__(self, exc_type, exc_value, exc_tb):
self._entered = False
if not exc_type and not exc_value and not exc_tb:
dst_buffer = ffi.new('char[]', self._write_size)
out_buffer = ffi.new('ZSTD_outBuffer *')
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer.dst = dst_buffer
out_buffer.size = len(dst_buffer)
out_buffer.pos = 0
in_buffer.src = ffi.NULL
in_buffer.size = 0
in_buffer.pos = 0
while True:
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
out_buffer, in_buffer,
lib.ZSTD_e_end)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
_zstd_error(zresult))
if out_buffer.pos:
self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
out_buffer.pos = 0
if zresult == 0:
break
self._compressor = None
return False
def memory_size(self):
if not self._entered:
raise ZstdError('cannot determine size of an inactive compressor; '
'call when a context manager is active')
return lib.ZSTD_sizeof_CCtx(self._compressor._cctx)
def write(self, data):
if not self._entered:
raise ZstdError('write() must be called from an active context '
'manager')
total_write = 0
data_buffer = ffi.from_buffer(data)
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
out_buffer = ffi.new('ZSTD_outBuffer *')
dst_buffer = ffi.new('char[]', self._write_size)
out_buffer.dst = dst_buffer
out_buffer.size = self._write_size
out_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
out_buffer, in_buffer,
lib.ZSTD_e_continue)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
total_write += out_buffer.pos
self._bytes_compressed += out_buffer.pos
out_buffer.pos = 0
return total_write
def flush(self):
if not self._entered:
raise ZstdError('flush must be called from an active context manager')
total_write = 0
out_buffer = ffi.new('ZSTD_outBuffer *')
dst_buffer = ffi.new('char[]', self._write_size)
out_buffer.dst = dst_buffer
out_buffer.size = self._write_size
out_buffer.pos = 0
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = ffi.NULL
in_buffer.size = 0
in_buffer.pos = 0
while True:
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
out_buffer, in_buffer,
lib.ZSTD_e_flush)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
if not out_buffer.pos:
break
self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
total_write += out_buffer.pos
self._bytes_compressed += out_buffer.pos
out_buffer.pos = 0
return total_write
def tell(self):
return self._bytes_compressed
class ZstdCompressionObj(object):
def compress(self, data):
if self._finished:
raise ZstdError('cannot call compress() after compressor finished')
data_buffer = ffi.from_buffer(data)
source = ffi.new('ZSTD_inBuffer *')
source.src = data_buffer
source.size = len(data_buffer)
source.pos = 0
chunks = []
while source.pos < len(data):
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
self._out,
source,
lib.ZSTD_e_continue)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
if self._out.pos:
chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])
self._out.pos = 0
return b''.join(chunks)
def flush(self, flush_mode=COMPRESSOBJ_FLUSH_FINISH):
if flush_mode not in (COMPRESSOBJ_FLUSH_FINISH, COMPRESSOBJ_FLUSH_BLOCK):
raise ValueError('flush mode not recognized')
if self._finished:
raise ZstdError('compressor object already finished')
assert self._out.pos == 0
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = ffi.NULL
in_buffer.size = 0
in_buffer.pos = 0
if flush_mode == COMPRESSOBJ_FLUSH_BLOCK:
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
self._out,
in_buffer,
lib.ZSTD_e_flush)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
# Output buffer is guaranteed to hold full block.
assert zresult == 0
if self._out.pos:
result = ffi.buffer(self._out.dst, self._out.pos)[:]
self._out.pos = 0
return result
else:
return b''
assert flush_mode == COMPRESSOBJ_FLUSH_FINISH
self._finished = True
chunks = []
while True:
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
self._out,
in_buffer,
lib.ZSTD_e_end)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
_zstd_error(zresult))
if self._out.pos:
chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])
self._out.pos = 0
if not zresult:
break
return b''.join(chunks)
class CompressionReader(object):
def __init__(self, compressor, source, size, read_size):
self._compressor = compressor
self._source = source
self._source_size = size
self._read_size = read_size
self._entered = False
self._closed = False
self._bytes_compressed = 0
self._finished_input = False
self._finished_output = False
self._in_buffer = ffi.new('ZSTD_inBuffer *')
# Holds a ref so backing bytes in self._in_buffer stay alive.
self._source_buffer = None
def __enter__(self):
if self._entered:
raise ValueError('cannot __enter__ multiple times')
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._compressor._cctx,
self._source_size)
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
self._entered = True
return self
def __exit__(self, exc_type, exc_value, exc_tb):
self._entered = False
self._closed = True
self._source = None
self._compressor = None
return False
def readable(self):
return True
def writable(self):
return False
def seekable(self):
return False
def readline(self):
raise io.UnsupportedOperation()
def readlines(self):
raise io.UnsupportedOperation()
def write(self, data):
raise OSError('stream is not writable')
def writelines(self, ignored):
raise OSError('stream is not writable')
def isatty(self):
return False
def flush(self):
return None
def close(self):
self._closed = True
return None
def closed(self):
return self._closed
def tell(self):
return self._bytes_compressed
def readall(self):
raise NotImplementedError()
def __iter__(self):
raise io.UnsupportedOperation()
def __next__(self):
raise io.UnsupportedOperation()
next = __next__
def read(self, size=-1):
if not self._entered:
raise ZstdError('read() must be called from an active context manager')
if self._closed:
raise ValueError('stream is closed')
if self._finished_output:
return b''
if size < 1:
raise ValueError('cannot read negative or size 0 amounts')
# Need a dedicated ref to dest buffer otherwise it gets collected.
dst_buffer = ffi.new('char[]', size)
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = dst_buffer
out_buffer.size = size
out_buffer.pos = 0
def compress_input():
if self._in_buffer.pos >= self._in_buffer.size:
return
old_pos = out_buffer.pos
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
out_buffer, self._in_buffer,
lib.ZSTD_e_continue)
self._bytes_compressed += out_buffer.pos - old_pos
if self._in_buffer.pos == self._in_buffer.size:
self._in_buffer.src = ffi.NULL
self._in_buffer.pos = 0
self._in_buffer.size = 0
self._source_buffer = None
if not hasattr(self._source, 'read'):
self._finished_input = True
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s',
_zstd_error(zresult))
if out_buffer.pos and out_buffer.pos == out_buffer.size:
return ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
def get_input():
if self._finished_input:
return
if hasattr(self._source, 'read'):
data = self._source.read(self._read_size)
if not data:
self._finished_input = True
return
self._source_buffer = ffi.from_buffer(data)
self._in_buffer.src = self._source_buffer
self._in_buffer.size = len(self._source_buffer)
self._in_buffer.pos = 0
else:
self._source_buffer = ffi.from_buffer(self._source)
self._in_buffer.src = self._source_buffer
self._in_buffer.size = len(self._source_buffer)
self._in_buffer.pos = 0
result = compress_input()
if result:
return result
while not self._finished_input:
get_input()
result = compress_input()
if result:
return result
# EOF
old_pos = out_buffer.pos
zresult = lib.ZSTD_compress_generic(self._compressor._cctx,
out_buffer, self._in_buffer,
lib.ZSTD_e_end)
self._bytes_compressed += out_buffer.pos - old_pos
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s',
_zstd_error(zresult))
if zresult == 0:
self._finished_output = True
return ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
class ZstdCompressor(object):
def __init__(self, level=3, dict_data=None, compression_params=None,
write_checksum=None, write_content_size=None,
write_dict_id=None, threads=0):
if level > lib.ZSTD_maxCLevel():
raise ValueError('level must be less than %d' % lib.ZSTD_maxCLevel())
if threads < 0:
threads = _cpu_count()
if compression_params and write_checksum is not None:
raise ValueError('cannot define compression_params and '
'write_checksum')
if compression_params and write_content_size is not None:
raise ValueError('cannot define compression_params and '
'write_content_size')
if compression_params and write_dict_id is not None:
raise ValueError('cannot define compression_params and '
'write_dict_id')
if compression_params and threads:
raise ValueError('cannot define compression_params and threads')
if compression_params:
self._params = _make_cctx_params(compression_params)
else:
if write_dict_id is None:
write_dict_id = True
params = lib.ZSTD_createCCtxParams()
if params == ffi.NULL:
raise MemoryError()
self._params = ffi.gc(params, lib.ZSTD_freeCCtxParams)
_set_compression_parameter(self._params,
lib.ZSTD_p_compressionLevel,
level)
_set_compression_parameter(
self._params,
lib.ZSTD_p_contentSizeFlag,
write_content_size if write_content_size is not None else 1)
_set_compression_parameter(self._params,
lib.ZSTD_p_checksumFlag,
1 if write_checksum else 0)
_set_compression_parameter(self._params,
lib.ZSTD_p_dictIDFlag,
1 if write_dict_id else 0)
if threads:
_set_compression_parameter(self._params,
lib.ZSTD_p_nbWorkers,
threads)
cctx = lib.ZSTD_createCCtx()
if cctx == ffi.NULL:
raise MemoryError()
self._cctx = cctx
self._dict_data = dict_data
# We defer setting up garbage collection until after calling
# _ensure_cctx() to ensure the memory size estimate is more accurate.
try:
self._ensure_cctx()
finally:
self._cctx = ffi.gc(cctx, lib.ZSTD_freeCCtx,
size=lib.ZSTD_sizeof_CCtx(cctx))
def _ensure_cctx(self):
lib.ZSTD_CCtx_reset(self._cctx)
zresult = lib.ZSTD_CCtx_setParametersUsingCCtxParams(self._cctx,
self._params)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not set compression parameters: %s' %
_zstd_error(zresult))
dict_data = self._dict_data
if dict_data:
if dict_data._cdict:
zresult = lib.ZSTD_CCtx_refCDict(self._cctx, dict_data._cdict)
else:
zresult = lib.ZSTD_CCtx_loadDictionary_advanced(
self._cctx, dict_data.as_bytes(), len(dict_data),
lib.ZSTD_dlm_byRef, dict_data._dict_type)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not load compression dictionary: %s' %
_zstd_error(zresult))
def memory_size(self):
return lib.ZSTD_sizeof_CCtx(self._cctx)
def compress(self, data):
self._ensure_cctx()
data_buffer = ffi.from_buffer(data)
dest_size = lib.ZSTD_compressBound(len(data_buffer))
out = new_nonzero('char[]', dest_size)
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, len(data_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
out_buffer = ffi.new('ZSTD_outBuffer *')
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer.dst = out
out_buffer.size = dest_size
out_buffer.pos = 0
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
zresult = lib.ZSTD_compress_generic(self._cctx,
out_buffer,
in_buffer,
lib.ZSTD_e_end)
if lib.ZSTD_isError(zresult):
raise ZstdError('cannot compress: %s' %
_zstd_error(zresult))
elif zresult:
raise ZstdError('unexpected partial frame flush')
return ffi.buffer(out, out_buffer.pos)[:]
def compressobj(self, size=-1):
self._ensure_cctx()
if size < 0:
size = lib.ZSTD_CONTENTSIZE_UNKNOWN
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
cobj = ZstdCompressionObj()
cobj._out = ffi.new('ZSTD_outBuffer *')
cobj._dst_buffer = ffi.new('char[]', COMPRESSION_RECOMMENDED_OUTPUT_SIZE)
cobj._out.dst = cobj._dst_buffer
cobj._out.size = COMPRESSION_RECOMMENDED_OUTPUT_SIZE
cobj._out.pos = 0
cobj._compressor = self
cobj._finished = False
return cobj
def copy_stream(self, ifh, ofh, size=-1,
read_size=COMPRESSION_RECOMMENDED_INPUT_SIZE,
write_size=COMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if not hasattr(ifh, 'read'):
raise ValueError('first argument must have a read() method')
if not hasattr(ofh, 'write'):
raise ValueError('second argument must have a write() method')
self._ensure_cctx()
if size < 0:
size = lib.ZSTD_CONTENTSIZE_UNKNOWN
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
dst_buffer = ffi.new('char[]', write_size)
out_buffer.dst = dst_buffer
out_buffer.size = write_size
out_buffer.pos = 0
total_read, total_write = 0, 0
while True:
data = ifh.read(read_size)
if not data:
break
data_buffer = ffi.from_buffer(data)
total_read += len(data_buffer)
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_compress_generic(self._cctx,
out_buffer,
in_buffer,
lib.ZSTD_e_continue)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
total_write += out_buffer.pos
out_buffer.pos = 0
# We've finished reading. Flush the compressor.
while True:
zresult = lib.ZSTD_compress_generic(self._cctx,
out_buffer,
in_buffer,
lib.ZSTD_e_end)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
_zstd_error(zresult))
if out_buffer.pos:
ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
total_write += out_buffer.pos
out_buffer.pos = 0
if zresult == 0:
break
return total_read, total_write
def stream_reader(self, source, size=-1,
read_size=COMPRESSION_RECOMMENDED_INPUT_SIZE):
self._ensure_cctx()
try:
size = len(source)
except Exception:
pass
if size < 0:
size = lib.ZSTD_CONTENTSIZE_UNKNOWN
return CompressionReader(self, source, size, read_size)
def stream_writer(self, writer, size=-1,
write_size=COMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if not hasattr(writer, 'write'):
raise ValueError('must pass an object with a write() method')
self._ensure_cctx()
if size < 0:
size = lib.ZSTD_CONTENTSIZE_UNKNOWN
return ZstdCompressionWriter(self, writer, size, write_size)
write_to = stream_writer
def read_to_iter(self, reader, size=-1,
read_size=COMPRESSION_RECOMMENDED_INPUT_SIZE,
write_size=COMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if hasattr(reader, 'read'):
have_read = True
elif hasattr(reader, '__getitem__'):
have_read = False
buffer_offset = 0
size = len(reader)
else:
raise ValueError('must pass an object with a read() method or '
'conforms to buffer protocol')
self._ensure_cctx()
if size < 0:
size = lib.ZSTD_CONTENTSIZE_UNKNOWN
zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
if lib.ZSTD_isError(zresult):
raise ZstdError('error setting source size: %s' %
_zstd_error(zresult))
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
in_buffer.src = ffi.NULL
in_buffer.size = 0
in_buffer.pos = 0
dst_buffer = ffi.new('char[]', write_size)
out_buffer.dst = dst_buffer
out_buffer.size = write_size
out_buffer.pos = 0
while True:
# We should never have output data sitting around after a previous
# iteration.
assert out_buffer.pos == 0
# Collect input data.
if have_read:
read_result = reader.read(read_size)
else:
remaining = len(reader) - buffer_offset
slice_size = min(remaining, read_size)
read_result = reader[buffer_offset:buffer_offset + slice_size]
buffer_offset += slice_size
# No new input data. Break out of the read loop.
if not read_result:
break
# Feed all read data into the compressor and emit output until
# exhausted.
read_buffer = ffi.from_buffer(read_result)
in_buffer.src = read_buffer
in_buffer.size = len(read_buffer)
in_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_compress_generic(self._cctx, out_buffer, in_buffer,
lib.ZSTD_e_continue)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
out_buffer.pos = 0
yield data
assert out_buffer.pos == 0
# And repeat the loop to collect more data.
continue
# If we get here, input is exhausted. End the stream and emit what
# remains.
while True:
assert out_buffer.pos == 0
zresult = lib.ZSTD_compress_generic(self._cctx,
out_buffer,
in_buffer,
lib.ZSTD_e_end)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
_zstd_error(zresult))
if out_buffer.pos:
data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
out_buffer.pos = 0
yield data
if zresult == 0:
break
read_from = read_to_iter
def frame_progression(self):
progression = lib.ZSTD_getFrameProgression(self._cctx)
return progression.ingested, progression.consumed, progression.produced
class FrameParameters(object):
def __init__(self, fparams):
self.content_size = fparams.frameContentSize
self.window_size = fparams.windowSize
self.dict_id = fparams.dictID
self.has_checksum = bool(fparams.checksumFlag)
def frame_content_size(data):
data_buffer = ffi.from_buffer(data)
size = lib.ZSTD_getFrameContentSize(data_buffer, len(data_buffer))
if size == lib.ZSTD_CONTENTSIZE_ERROR:
raise ZstdError('error when determining content size')
elif size == lib.ZSTD_CONTENTSIZE_UNKNOWN:
return -1
else:
return size
def frame_header_size(data):
data_buffer = ffi.from_buffer(data)
zresult = lib.ZSTD_frameHeaderSize(data_buffer, len(data_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('could not determine frame header size: %s' %
_zstd_error(zresult))
return zresult
def get_frame_parameters(data):
params = ffi.new('ZSTD_frameHeader *')
data_buffer = ffi.from_buffer(data)
zresult = lib.ZSTD_getFrameHeader(params, data_buffer, len(data_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('cannot get frame parameters: %s' %
_zstd_error(zresult))
if zresult:
raise ZstdError('not enough data for frame parameters; need %d bytes' %
zresult)
return FrameParameters(params[0])
class ZstdCompressionDict(object):
def __init__(self, data, dict_type=DICT_TYPE_AUTO, k=0, d=0):
assert isinstance(data, bytes_type)
self._data = data
self.k = k
self.d = d
if dict_type not in (DICT_TYPE_AUTO, DICT_TYPE_RAWCONTENT,
DICT_TYPE_FULLDICT):
raise ValueError('invalid dictionary load mode: %d; must use '
'DICT_TYPE_* constants')
self._dict_type = dict_type
self._cdict = None
def __len__(self):
return len(self._data)
def dict_id(self):
return int_type(lib.ZDICT_getDictID(self._data, len(self._data)))
def as_bytes(self):
return self._data
def precompute_compress(self, level=0, compression_params=None):
if level and compression_params:
raise ValueError('must only specify one of level or '
'compression_params')
if not level and not compression_params:
raise ValueError('must specify one of level or compression_params')
if level:
cparams = lib.ZSTD_getCParams(level, 0, len(self._data))
else:
cparams = ffi.new('ZSTD_compressionParameters')
cparams.chainLog = compression_params.chain_log
cparams.hashLog = compression_params.hash_log
cparams.searchLength = compression_params.min_match
cparams.searchLog = compression_params.search_log
cparams.strategy = compression_params.compression_strategy
cparams.targetLength = compression_params.target_length
cparams.windowLog = compression_params.window_log
cdict = lib.ZSTD_createCDict_advanced(self._data, len(self._data),
lib.ZSTD_dlm_byRef,
self._dict_type,
cparams,
lib.ZSTD_defaultCMem)
if cdict == ffi.NULL:
raise ZstdError('unable to precompute dictionary')
self._cdict = ffi.gc(cdict, lib.ZSTD_freeCDict,
size=lib.ZSTD_sizeof_CDict(cdict))
@property
def _ddict(self):
ddict = lib.ZSTD_createDDict_advanced(self._data, len(self._data),
lib.ZSTD_dlm_byRef,
self._dict_type,
lib.ZSTD_defaultCMem)
if ddict == ffi.NULL:
raise ZstdError('could not create decompression dict')
ddict = ffi.gc(ddict, lib.ZSTD_freeDDict,
size=lib.ZSTD_sizeof_DDict(ddict))
self.__dict__['_ddict'] = ddict
return ddict
def train_dictionary(dict_size, samples, k=0, d=0, notifications=0, dict_id=0,
level=0, steps=0, threads=0):
if not isinstance(samples, list):
raise TypeError('samples must be a list')
if threads < 0:
threads = _cpu_count()
total_size = sum(map(len, samples))
samples_buffer = new_nonzero('char[]', total_size)
sample_sizes = new_nonzero('size_t[]', len(samples))
offset = 0
for i, sample in enumerate(samples):
if not isinstance(sample, bytes_type):
raise ValueError('samples must be bytes')
l = len(sample)
ffi.memmove(samples_buffer + offset, sample, l)
offset += l
sample_sizes[i] = l
dict_data = new_nonzero('char[]', dict_size)
dparams = ffi.new('ZDICT_cover_params_t *')[0]
dparams.k = k
dparams.d = d
dparams.steps = steps
dparams.nbThreads = threads
dparams.zParams.notificationLevel = notifications
dparams.zParams.dictID = dict_id
dparams.zParams.compressionLevel = level
if (not dparams.k and not dparams.d and not dparams.steps
and not dparams.nbThreads and not dparams.zParams.notificationLevel
and not dparams.zParams.dictID
and not dparams.zParams.compressionLevel):
zresult = lib.ZDICT_trainFromBuffer(
ffi.addressof(dict_data), dict_size,
ffi.addressof(samples_buffer),
ffi.addressof(sample_sizes, 0), len(samples))
elif dparams.steps or dparams.nbThreads:
zresult = lib.ZDICT_optimizeTrainFromBuffer_cover(
ffi.addressof(dict_data), dict_size,
ffi.addressof(samples_buffer),
ffi.addressof(sample_sizes, 0), len(samples),
ffi.addressof(dparams))
else:
zresult = lib.ZDICT_trainFromBuffer_cover(
ffi.addressof(dict_data), dict_size,
ffi.addressof(samples_buffer),
ffi.addressof(sample_sizes, 0), len(samples),
dparams)
if lib.ZDICT_isError(zresult):
msg = ffi.string(lib.ZDICT_getErrorName(zresult)).decode('utf-8')
raise ZstdError('cannot train dict: %s' % msg)
return ZstdCompressionDict(ffi.buffer(dict_data, zresult)[:],
dict_type=DICT_TYPE_FULLDICT,
k=dparams.k, d=dparams.d)
class ZstdDecompressionObj(object):
def __init__(self, decompressor, write_size):
self._decompressor = decompressor
self._write_size = write_size
self._finished = False
def decompress(self, data):
if self._finished:
raise ZstdError('cannot use a decompressobj multiple times')
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
data_buffer = ffi.from_buffer(data)
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
dst_buffer = ffi.new('char[]', self._write_size)
out_buffer.dst = dst_buffer
out_buffer.size = len(dst_buffer)
out_buffer.pos = 0
chunks = []
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_decompress_generic(self._decompressor._dctx,
out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompressor error: %s' %
_zstd_error(zresult))
if zresult == 0:
self._finished = True
self._decompressor = None
if out_buffer.pos:
chunks.append(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
out_buffer.pos = 0
return b''.join(chunks)
class DecompressionReader(object):
def __init__(self, decompressor, source, read_size):
self._decompressor = decompressor
self._source = source
self._read_size = read_size
self._entered = False
self._closed = False
self._bytes_decompressed = 0
self._finished_input = False
self._finished_output = False
self._in_buffer = ffi.new('ZSTD_inBuffer *')
# Holds a ref to self._in_buffer.src.
self._source_buffer = None
def __enter__(self):
if self._entered:
raise ValueError('cannot __enter__ multiple times')
self._decompressor._ensure_dctx()
self._entered = True
return self
def __exit__(self, exc_type, exc_value, exc_tb):
self._entered = False
self._closed = True
self._source = None
self._decompressor = None
return False
def readable(self):
return True
def writable(self):
return False
def seekable(self):
return True
def readline(self):
raise NotImplementedError()
def readlines(self):
raise NotImplementedError()
def write(self, data):
raise io.UnsupportedOperation()
def writelines(self, lines):
raise io.UnsupportedOperation()
def isatty(self):
return False
def flush(self):
return None
def close(self):
self._closed = True
return None
def closed(self):
return self._closed
def tell(self):
return self._bytes_decompressed
def readall(self):
raise NotImplementedError()
def __iter__(self):
raise NotImplementedError()
def __next__(self):
raise NotImplementedError()
next = __next__
def read(self, size=-1):
if not self._entered:
raise ZstdError('read() must be called from an active context manager')
if self._closed:
raise ValueError('stream is closed')
if self._finished_output:
return b''
if size < 1:
raise ValueError('cannot read negative or size 0 amounts')
dst_buffer = ffi.new('char[]', size)
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = dst_buffer
out_buffer.size = size
out_buffer.pos = 0
def decompress():
zresult = lib.ZSTD_decompress_generic(self._decompressor._dctx,
out_buffer, self._in_buffer)
if self._in_buffer.pos == self._in_buffer.size:
self._in_buffer.src = ffi.NULL
self._in_buffer.pos = 0
self._in_buffer.size = 0
self._source_buffer = None
if not hasattr(self._source, 'read'):
self._finished_input = True
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompress error: %s',
_zstd_error(zresult))
elif zresult == 0:
self._finished_output = True
if out_buffer.pos and out_buffer.pos == out_buffer.size:
self._bytes_decompressed += out_buffer.size
return ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
def get_input():
if self._finished_input:
return
if hasattr(self._source, 'read'):
data = self._source.read(self._read_size)
if not data:
self._finished_input = True
return
self._source_buffer = ffi.from_buffer(data)
self._in_buffer.src = self._source_buffer
self._in_buffer.size = len(self._source_buffer)
self._in_buffer.pos = 0
else:
self._source_buffer = ffi.from_buffer(self._source)
self._in_buffer.src = self._source_buffer
self._in_buffer.size = len(self._source_buffer)
self._in_buffer.pos = 0
get_input()
result = decompress()
if result:
return result
while not self._finished_input:
get_input()
result = decompress()
if result:
return result
self._bytes_decompressed += out_buffer.pos
return ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
def seek(self, pos, whence=os.SEEK_SET):
if not self._entered:
raise ZstdError('seek() must be called from an active context '
'manager')
if self._closed:
raise ValueError('stream is closed')
read_amount = 0
if whence == os.SEEK_SET:
if pos < 0:
raise ValueError('cannot seek to negative position with SEEK_SET')
if pos < self._bytes_decompressed:
raise ValueError('cannot seek zstd decompression stream '
'backwards')
read_amount = pos - self._bytes_decompressed
elif whence == os.SEEK_CUR:
if pos < 0:
raise ValueError('cannot seek zstd decompression stream '
'backwards')
read_amount = pos
elif whence == os.SEEK_END:
raise ValueError('zstd decompression streams cannot be seeked '
'with SEEK_END')
while read_amount:
result = self.read(min(read_amount,
DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE))
if not result:
break
read_amount -= len(result)
return self._bytes_decompressed
class ZstdDecompressionWriter(object):
def __init__(self, decompressor, writer, write_size):
self._decompressor = decompressor
self._writer = writer
self._write_size = write_size
self._entered = False
def __enter__(self):
if self._entered:
raise ZstdError('cannot __enter__ multiple times')
self._decompressor._ensure_dctx()
self._entered = True
return self
def __exit__(self, exc_type, exc_value, exc_tb):
self._entered = False
def memory_size(self):
if not self._decompressor._dctx:
raise ZstdError('cannot determine size of inactive decompressor '
'call when context manager is active')
return lib.ZSTD_sizeof_DCtx(self._decompressor._dctx)
def write(self, data):
if not self._entered:
raise ZstdError('write must be called from an active context manager')
total_write = 0
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
data_buffer = ffi.from_buffer(data)
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
dst_buffer = ffi.new('char[]', self._write_size)
out_buffer.dst = dst_buffer
out_buffer.size = len(dst_buffer)
out_buffer.pos = 0
dctx = self._decompressor._dctx
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_decompress_generic(dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompress error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
total_write += out_buffer.pos
out_buffer.pos = 0
return total_write
class ZstdDecompressor(object):
def __init__(self, dict_data=None, max_window_size=0, format=FORMAT_ZSTD1):
self._dict_data = dict_data
self._max_window_size = max_window_size
self._format = format
dctx = lib.ZSTD_createDCtx()
if dctx == ffi.NULL:
raise MemoryError()
self._dctx = dctx
# Defer setting up garbage collection until full state is loaded so
# the memory size is more accurate.
try:
self._ensure_dctx()
finally:
self._dctx = ffi.gc(dctx, lib.ZSTD_freeDCtx,
size=lib.ZSTD_sizeof_DCtx(dctx))
def memory_size(self):
return lib.ZSTD_sizeof_DCtx(self._dctx)
def decompress(self, data, max_output_size=0):
self._ensure_dctx()
data_buffer = ffi.from_buffer(data)
output_size = lib.ZSTD_getFrameContentSize(data_buffer, len(data_buffer))
if output_size == lib.ZSTD_CONTENTSIZE_ERROR:
raise ZstdError('error determining content size from frame header')
elif output_size == 0:
return b''
elif output_size == lib.ZSTD_CONTENTSIZE_UNKNOWN:
if not max_output_size:
raise ZstdError('could not determine content size in frame header')
result_buffer = ffi.new('char[]', max_output_size)
result_size = max_output_size
output_size = 0
else:
result_buffer = ffi.new('char[]', output_size)
result_size = output_size
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = result_buffer
out_buffer.size = result_size
out_buffer.pos = 0
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
zresult = lib.ZSTD_decompress_generic(self._dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('decompression error: %s' %
_zstd_error(zresult))
elif zresult:
raise ZstdError('decompression error: did not decompress full frame')
elif output_size and out_buffer.pos != output_size:
raise ZstdError('decompression error: decompressed %d bytes; expected %d' %
(zresult, output_size))
return ffi.buffer(result_buffer, out_buffer.pos)[:]
def stream_reader(self, source, read_size=DECOMPRESSION_RECOMMENDED_INPUT_SIZE):
self._ensure_dctx()
return DecompressionReader(self, source, read_size)
def decompressobj(self, write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if write_size < 1:
raise ValueError('write_size must be positive')
self._ensure_dctx()
return ZstdDecompressionObj(self, write_size=write_size)
def read_to_iter(self, reader, read_size=DECOMPRESSION_RECOMMENDED_INPUT_SIZE,
write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE,
skip_bytes=0):
if skip_bytes >= read_size:
raise ValueError('skip_bytes must be smaller than read_size')
if hasattr(reader, 'read'):
have_read = True
elif hasattr(reader, '__getitem__'):
have_read = False
buffer_offset = 0
size = len(reader)
else:
raise ValueError('must pass an object with a read() method or '
'conforms to buffer protocol')
if skip_bytes:
if have_read:
reader.read(skip_bytes)
else:
if skip_bytes > size:
raise ValueError('skip_bytes larger than first input chunk')
buffer_offset = skip_bytes
self._ensure_dctx()
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
dst_buffer = ffi.new('char[]', write_size)
out_buffer.dst = dst_buffer
out_buffer.size = len(dst_buffer)
out_buffer.pos = 0
while True:
assert out_buffer.pos == 0
if have_read:
read_result = reader.read(read_size)
else:
remaining = size - buffer_offset
slice_size = min(remaining, read_size)
read_result = reader[buffer_offset:buffer_offset + slice_size]
buffer_offset += slice_size
# No new input. Break out of read loop.
if not read_result:
break
# Feed all read data into decompressor and emit output until
# exhausted.
read_buffer = ffi.from_buffer(read_result)
in_buffer.src = read_buffer
in_buffer.size = len(read_buffer)
in_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
assert out_buffer.pos == 0
zresult = lib.ZSTD_decompress_generic(self._dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompress error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
out_buffer.pos = 0
yield data
if zresult == 0:
return
# Repeat loop to collect more input data.
continue
# If we get here, input is exhausted.
read_from = read_to_iter
def stream_writer(self, writer, write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if not hasattr(writer, 'write'):
raise ValueError('must pass an object with a write() method')
return ZstdDecompressionWriter(self, writer, write_size)
write_to = stream_writer
def copy_stream(self, ifh, ofh,
read_size=DECOMPRESSION_RECOMMENDED_INPUT_SIZE,
write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE):
if not hasattr(ifh, 'read'):
raise ValueError('first argument must have a read() method')
if not hasattr(ofh, 'write'):
raise ValueError('second argument must have a write() method')
self._ensure_dctx()
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
dst_buffer = ffi.new('char[]', write_size)
out_buffer.dst = dst_buffer
out_buffer.size = write_size
out_buffer.pos = 0
total_read, total_write = 0, 0
# Read all available input.
while True:
data = ifh.read(read_size)
if not data:
break
data_buffer = ffi.from_buffer(data)
total_read += len(data_buffer)
in_buffer.src = data_buffer
in_buffer.size = len(data_buffer)
in_buffer.pos = 0
# Flush all read data to output.
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_decompress_generic(self._dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompressor error: %s' %
_zstd_error(zresult))
if out_buffer.pos:
ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
total_write += out_buffer.pos
out_buffer.pos = 0
# Continue loop to keep reading.
return total_read, total_write
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_frameHeader *')
zresult = lib.ZSTD_getFrameHeader(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 params.frameContentSize == lib.ZSTD_CONTENTSIZE_UNKNOWN:
raise ValueError('chunk 0 missing content size in frame')
self._ensure_dctx(load_dict=False)
last_buffer = ffi.new('char[]', params.frameContentSize)
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = last_buffer
out_buffer.size = len(last_buffer)
out_buffer.pos = 0
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = chunk_buffer
in_buffer.size = len(chunk_buffer)
in_buffer.pos = 0
zresult = lib.ZSTD_decompress_generic(self._dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not decompress chunk 0: %s' %
_zstd_error(zresult))
elif zresult:
raise ZstdError('chunk 0 did not decompress full frame')
# 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_getFrameHeader(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 params.frameContentSize == lib.ZSTD_CONTENTSIZE_UNKNOWN:
raise ValueError('chunk %d missing content size in frame' % i)
dest_buffer = ffi.new('char[]', params.frameContentSize)
out_buffer.dst = dest_buffer
out_buffer.size = len(dest_buffer)
out_buffer.pos = 0
in_buffer.src = chunk_buffer
in_buffer.size = len(chunk_buffer)
in_buffer.pos = 0
zresult = lib.ZSTD_decompress_generic(self._dctx, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('could not decompress chunk %d: %s' %
_zstd_error(zresult))
elif zresult:
raise ZstdError('chunk %d did not decompress full frame' % i)
last_buffer = dest_buffer
i += 1
return ffi.buffer(last_buffer, len(last_buffer))[:]
def _ensure_dctx(self, load_dict=True):
lib.ZSTD_DCtx_reset(self._dctx)
if self._max_window_size:
zresult = lib.ZSTD_DCtx_setMaxWindowSize(self._dctx,
self._max_window_size)
if lib.ZSTD_isError(zresult):
raise ZstdError('unable to set max window size: %s' %
_zstd_error(zresult))
zresult = lib.ZSTD_DCtx_setFormat(self._dctx, self._format)
if lib.ZSTD_isError(zresult):
raise ZstdError('unable to set decoding format: %s' %
_zstd_error(zresult))
if self._dict_data and load_dict:
zresult = lib.ZSTD_DCtx_refDDict(self._dctx, self._dict_data._ddict)
if lib.ZSTD_isError(zresult):
raise ZstdError('unable to reference prepared dictionary: %s' %
_zstd_error(zresult))