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 copy_stream(self, ifh, ofh):
cstream = self._get_cstream()
in_buffer = ffi.new('ZSTD_inBuffer *')
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = ffi.new('char[]', _CSTREAM_OUT_SIZE)
out_buffer.size = _CSTREAM_OUT_SIZE
out_buffer.pos = 0
total_read, total_write = 0, 0
while True:
data = ifh.read(_CSTREAM_IN_SIZE)
if not data:
break
total_read += len(data)
in_buffer.src = ffi.new('char[]', data)
in_buffer.size = len(data)
in_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
res = lib.ZSTD_compressStream(cstream, out_buffer, in_buffer)
if lib.ZSTD_isError(res):
raise Exception('zstd compress error: %s' %
lib.ZSTD_getErrorName(res))
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:
res = lib.ZSTD_endStream(cstream, out_buffer)
if lib.ZSTD_isError(res):
raise Exception('error ending compression stream: %s' %
lib.ZSTD_getErrorName(res))
if out_buffer.pos:
ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
total_write += out_buffer.pos
out_buffer.pos = 0
if res == 0:
break
return total_read, total_write
def write(self, data):
out_buffer = ffi.new('ZSTD_outBuffer *')
out_buffer.dst = ffi.new('char[]', _CSTREAM_OUT_SIZE)
out_buffer.size = _CSTREAM_OUT_SIZE
out_buffer.pos = 0
# TODO can we reuse existing memory?
in_buffer = ffi.new('ZSTD_inBuffer *')
in_buffer.src = ffi.new('char[]', data)
in_buffer.size = len(data)
in_buffer.pos = 0
while in_buffer.pos < in_buffer.size:
res = lib.ZSTD_compressStream(self._cstream, out_buffer, in_buffer)
if lib.ZSTD_isError(res):
raise Exception('zstd compress error: %s' % lib.ZSTD_getErrorName(res))
if out_buffer.pos:
self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
out_buffer.pos = 0
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:
if self._mtcctx:
zresult = lib.ZSTDMT_compressStream(self._mtcctx, out_buffer,
in_buffer)
else:
zresult = lib.ZSTD_compressStream(self._compressor._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:
self._writer.write(
ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
total_write += out_buffer.pos
out_buffer.pos = 0
return total_write
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