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):
if self._mtcctx:
zresult = lib.ZSTDMT_compressStream(self._mtcctx, self._out,
source)
else:
zresult = lib.ZSTD_compressStream(self._compressor._cstream,
self._out, source)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
if self._out.pos:
chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])
self._out.pos = 0
return b''.join(chunks)
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
dstream = self._decompressor._dstream
while in_buffer.pos < in_buffer.size:
zresult = lib.ZSTD_decompressStream(dstream, out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompress error: %s' %
ffi.string(lib.ZSTD_getErrorName(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
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_dstream()
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_decompressStream(self._dstream, out_buffer,
in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompressor error: %s' %
ffi.string(lib.ZSTD_getErrorName(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(self, data, max_output_size=0):
data_buffer = ffi.from_buffer(data)
orig_dctx = new_nonzero('char[]', lib.ZSTD_sizeof_DCtx(self._refdctx))
dctx = ffi.cast('ZSTD_DCtx *', orig_dctx)
lib.ZSTD_copyDCtx(dctx, self._refdctx)
ddict = self._ddict
output_size = lib.ZSTD_getDecompressedSize(data_buffer,
len(data_buffer))
if output_size:
result_buffer = ffi.new('char[]', output_size)
result_size = output_size
else:
if not max_output_size:
raise ZstdError('input data invalid or missing content size '
'in frame header')
result_buffer = ffi.new('char[]', max_output_size)
result_size = max_output_size
if ddict:
zresult = lib.ZSTD_decompress_usingDDict(dctx, result_buffer,
result_size, data_buffer,
len(data_buffer), ddict)
else:
zresult = lib.ZSTD_decompressDCtx(dctx, result_buffer, result_size,
data_buffer, len(data_buffer))
if lib.ZSTD_isError(zresult):
raise ZstdError('decompression error: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
elif output_size and zresult != output_size:
raise ZstdError(
'decompression error: decompressed %d bytes; expected %d' %
(zresult, output_size))
return ffi.buffer(result_buffer, zresult)[:]
def decompress(self, data):
if self._finished:
raise ZstdError('cannot use a decompressobj multiple times')
assert (self._decompressor._dstream)
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[]', DECOMPRESSION_RECOMMENDED_OUTPUT_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_decompressStream(self._decompressor._dstream,
out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompressor error: %s' %
ffi.string(lib.ZSTD_getErrorName(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)
def read_from(self,
reader,
size=0,
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')
if self._multithreaded:
self._init_mtcstream(size)
else:
self._ensure_cstream(size)
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:
if self._multithreaded:
zresult = lib.ZSTDMT_compressStream(
self._cctx, out_buffer, in_buffer)
else:
zresult = lib.ZSTD_compressStream(self._cstream,
out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
ffi.string(lib.ZSTD_getErrorName(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
if self._multithreaded:
zresult = lib.ZSTDMT_endStream(self._cctx, out_buffer)
else:
zresult = lib.ZSTD_endStream(self._cstream, out_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
ffi.string(lib.ZSTD_getErrorName(zresult)))
if out_buffer.pos:
data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
out_buffer.pos = 0
yield data
if zresult == 0:
break
def copy_stream(self,
ifh,
ofh,
size=0,
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')
mt = self._multithreaded
if mt:
self._init_mtcstream(size)
else:
self._ensure_cstream(size)
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:
if mt:
zresult = lib.ZSTDMT_compressStream(
self._cctx, out_buffer, in_buffer)
else:
zresult = lib.ZSTD_compressStream(self._cstream,
out_buffer, in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd compress error: %s' %
ffi.string(lib.ZSTD_getErrorName(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:
if mt:
zresult = lib.ZSTDMT_endStream(self._cctx, out_buffer)
else:
zresult = lib.ZSTD_endStream(self._cstream, out_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('error ending compression stream: %s' %
ffi.string(lib.ZSTD_getErrorName(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 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))[:]
def read_from(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_dstream()
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_decompressStream(self._dstream, out_buffer,
in_buffer)
if lib.ZSTD_isError(zresult):
raise ZstdError('zstd decompress error: %s' %
ffi.string(lib.ZSTD_getErrorName(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