Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
ValueError, "chunk 0 is too small to contain a zstd frame"
):
dctx.decompress_content_dict_chain([zstd.FRAME_HEADER])
with self.assertRaisesRegex(ValueError, "chunk 0 is not a valid zstd frame"):
dctx.decompress_content_dict_chain([b"foo" * 8])
no_size = zstd.ZstdCompressor(write_content_size=False).compress(b"foo" * 64)
with self.assertRaisesRegex(
ValueError, "chunk 0 missing content size in frame"
):
dctx.decompress_content_dict_chain([no_size])
# Corrupt first frame.
frame = zstd.ZstdCompressor().compress(b"foo" * 64)
frame = frame[0:12] + frame[15:]
with self.assertRaisesRegex(
zstd.ZstdError, "chunk 0 did not decompress full frame"
):
dctx.decompress_content_dict_chain([frame])
def test_data_equivalence(self, original, threads, use_dict):
kwargs = {}
# Use a content dictionary because it is cheap to create.
if use_dict:
kwargs["dict_data"] = zstd.ZstdCompressionDict(original[0])
cctx = zstd.ZstdCompressor(level=1, write_checksum=True, **kwargs)
if not hasattr(cctx, "multi_compress_to_buffer"):
self.skipTest("multi_compress_to_buffer not available")
result = cctx.multi_compress_to_buffer(original, threads=-1)
self.assertEqual(len(result), len(original))
# The frame produced via the batch APIs may not be bit identical to that
# produced by compress() because compression parameters are adjusted
# from the first input in batch mode. So the only thing we can do is
# verify the decompressed data matches the input.
dctx = zstd.ZstdDecompressor(**kwargs)
for i, frame in enumerate(result):
self.assertEqual(dctx.decompress(frame), original[i])
def test_multiple_threads(self):
# threads argument will cause multi-threaded ZSTD APIs to be used, which will
# make output different.
refcctx = zstd.ZstdCompressor(write_checksum=True)
reference = [refcctx.compress(b"x" * 64), refcctx.compress(b"y" * 64)]
cctx = zstd.ZstdCompressor(write_checksum=True)
if not hasattr(cctx, "multi_compress_to_buffer"):
self.skipTest("multi_compress_to_buffer not available")
frames = []
frames.extend(b"x" * 64 for i in range(256))
frames.extend(b"y" * 64 for i in range(256))
result = cctx.multi_compress_to_buffer(frames, threads=-1)
self.assertEqual(len(result), 512)
for i in range(512):
if i < 256:
self.assertEqual(result[i].tobytes(), reference[0])
else:
self.assertEqual(result[i].tobytes(), reference[1])
def test_write_size(self):
cctx = zstd.ZstdCompressor(level=3)
dest = OpCountingBytesIO()
with cctx.stream_writer(dest, write_size=1) as compressor:
self.assertEqual(compressor.write(b"foo"), 0)
self.assertEqual(compressor.write(b"bar"), 0)
self.assertEqual(compressor.write(b"foobar"), 0)
self.assertEqual(len(dest.getvalue()), dest._write_count)
def test_stream_source_read_variance(
self, original, level, source_read_size, read_sizes
):
refctx = zstd.ZstdCompressor(level=level)
ref_frame = refctx.compress(original)
cctx = zstd.ZstdCompressor(level=level)
with cctx.stream_reader(
io.BytesIO(original), size=len(original), read_size=source_read_size
) as reader:
chunks = []
while True:
read_size = read_sizes.draw(strategies.integers(-1, 16384))
chunk = reader.read(read_size)
if not chunk and read_size:
break
chunks.append(chunk)
self.assertEqual(b"".join(chunks), ref_frame)
def test_no_context_manager(self):
source = b"foobar" * 60
cctx = zstd.ZstdCompressor()
frame = cctx.compress(source)
dctx = zstd.ZstdDecompressor()
reader = dctx.stream_reader(frame)
self.assertEqual(reader.read(6), b"foobar")
self.assertEqual(reader.read(18), b"foobar" * 3)
self.assertFalse(reader.closed)
# Calling close prevents subsequent use.
reader.close()
self.assertTrue(reader.closed)
with self.assertRaisesRegex(ValueError, "stream is closed"):
reader.read(6)
def test_relative_seeks(
self, original, level, source_read_size, seek_amounts, read_sizes
):
cctx = zstd.ZstdCompressor(level=level)
frame = cctx.compress(original)
dctx = zstd.ZstdDecompressor()
with dctx.stream_reader(frame, read_size=source_read_size) as reader:
while True:
amount = seek_amounts.draw(strategies.integers(0, 16384))
reader.seek(amount, os.SEEK_CUR)
offset = reader.tell()
read_amount = read_sizes.draw(strategies.integers(1, 16384))
chunk = reader.read(read_amount)
if not chunk:
break
def compress_content_dict_read_to_iter(chunks, zparams, use_size=False):
zctx = zstd.ZstdCompressor(compression_params=zparams)
size = len(chunks[0]) if use_size else -1
for o in zctx.read_to_iter(chunks[0], size=size):
pass
for i, chunk in enumerate(chunks[1:]):
d = zstd.ZstdCompressionDict(chunks[i])
zctx = zstd.ZstdCompressor(dict_data=d, compression_params=zparams)
size = len(chunk) if use_size else -1
for o in zctx.read_to_iter(chunk, size=size):
pass
def compress_content_dict_read_to_iter(chunks, zparams, use_size=False):
zctx = zstd.ZstdCompressor(compression_params=zparams)
size = len(chunks[0]) if use_size else -1
for o in zctx.read_to_iter(chunks[0], size=size):
pass
for i, chunk in enumerate(chunks[1:]):
d = zstd.ZstdCompressionDict(chunks[i])
zctx = zstd.ZstdCompressor(dict_data=d, compression_params=zparams)
size = len(chunk) if use_size else -1
for o in zctx.read_to_iter(chunk, size=size):
pass
def add(self, src, name=None, stop_event=None):
name = name or os.path.basename(src)
self.log(logging.DEBUG, "Add %s into %s", src, self.archive_path(name))
cctx = zstd.ZstdCompressor(compression_params=self.zstd_params)
try:
with open(src, "rb") as ifh, open(self.archive_path(name), "wb") as ofh:
with cctx.stream_writer(ofh) as writer:
while True:
if stop_event and stop_event.is_set():
raise CancelledError()
data = ifh.read(zstd.COMPRESSION_RECOMMENDED_INPUT_SIZE)
if not data:
break
if stop_event and stop_event.is_set():
raise CancelledError()
writer.write(data)
except:
if os.path.exists(self.archive_path(name)):