def test_buffer_with_segments_collection_input(self):
cctx = zstd.ZstdCompressor(write_content_size=True,
write_checksum=True)
original = [
b'foo1',
b'foo2' * 2,
b'foo3' * 3,
b'foo4' * 4,
b'foo5' * 5,
]
frames = [cctx.compress(c) for c in original]
b = b''.join([original[0], original[1]])
b1 = zstd.BufferWithSegments(
b,
struct.pack('=QQQQ', 0, len(original[0]), len(original[0]),
len(original[1])))
b = b''.join([original[2], original[3], original[4]])
b2 = zstd.BufferWithSegments(
b,
struct.pack('=QQQQQQ', 0, len(original[2]), len(original[2]),
len(original[3]),
len(original[2]) + len(original[3]), len(original[4])))
c = zstd.BufferWithSegmentsCollection(b1, b2)
result = cctx.multi_compress_to_buffer(c)
self.assertEqual(len(result), len(frames))
for i, frame in enumerate(frames):
self.assertEqual(result[i].tobytes(), frame)
def test_arguments(self):
with self.assertRaises(TypeError):
zstd.BufferWithSegments()
with self.assertRaises(TypeError):
zstd.BufferWithSegments(b'foo')
# Segments data should be a multiple of 16.
with self.assertRaisesRegexp(
ValueError, 'segments array size is not a multiple of 16'):
zstd.BufferWithSegments(b'foo', b'\x00\x00')
def test_argument_validation(self):
with self.assertRaisesRegexp(TypeError,
'arguments must be BufferWithSegments'):
zstd.BufferWithSegmentsCollection(None)
with self.assertRaisesRegexp(TypeError,
'arguments must be BufferWithSegments'):
zstd.BufferWithSegmentsCollection(
zstd.BufferWithSegments(b'foo', ss.pack(0, 3)), None)
with self.assertRaisesRegexp(ValueError,
'ZstdBufferWithSegments cannot be empty'):
zstd.BufferWithSegmentsCollection(zstd.BufferWithSegments(
b'', b''))
def test_length(self):
b1 = zstd.BufferWithSegments(b'foo', ss.pack(0, 3))
b2 = zstd.BufferWithSegments(b'barbaz',
b''.join([ss.pack(0, 3),
ss.pack(3, 3)]))
c = zstd.BufferWithSegmentsCollection(b1)
self.assertEqual(len(c), 1)
self.assertEqual(c.size(), 3)
c = zstd.BufferWithSegmentsCollection(b2)
self.assertEqual(len(c), 2)
self.assertEqual(c.size(), 6)
c = zstd.BufferWithSegmentsCollection(b1, b2)
self.assertEqual(len(c), 3)
self.assertEqual(c.size(), 9)
def test_getitem(self):
b1 = zstd.BufferWithSegments(b'foo', ss.pack(0, 3))
b2 = zstd.BufferWithSegments(b'barbaz',
b''.join([ss.pack(0, 3),
ss.pack(3, 3)]))
c = zstd.BufferWithSegmentsCollection(b1, b2)
with self.assertRaisesRegexp(IndexError, 'offset must be less than 3'):
c[3]
with self.assertRaisesRegexp(IndexError, 'offset must be less than 3'):
c[4]
self.assertEqual(c[0].tobytes(), b'foo')
self.assertEqual(c[1].tobytes(), b'bar')
self.assertEqual(c[2].tobytes(), b'baz')
def test_buffer_with_segments_collection_input(self):
cctx = zstd.ZstdCompressor(write_content_size=True)
original = [
b'foo0' * 2,
b'foo1' * 3,
b'foo2' * 4,
b'foo3' * 5,
b'foo4' * 6,
]
frames = cctx.multi_compress_to_buffer(original)
# Check round trip.
dctx = zstd.ZstdDecompressor()
decompressed = dctx.multi_decompress_to_buffer(frames, threads=3)
self.assertEqual(len(decompressed), len(original))
for i, data in enumerate(original):
self.assertEqual(data, decompressed[i].tobytes())
# And a manual mode.
b = b''.join([frames[0].tobytes(), frames[1].tobytes()])
b1 = zstd.BufferWithSegments(
b,
struct.pack('=QQQQ', 0, len(frames[0]), len(frames[0]),
len(frames[1])))
b = b''.join(
[frames[2].tobytes(), frames[3].tobytes(), frames[4].tobytes()])
b2 = zstd.BufferWithSegments(
b,
struct.pack('=QQQQQQ', 0, len(frames[2]), len(frames[2]),
len(frames[3]),
len(frames[2]) + len(frames[3]), len(frames[4])))
c = zstd.BufferWithSegmentsCollection(b1, b2)
dctx = zstd.ZstdDecompressor()
decompressed = dctx.multi_decompress_to_buffer(c)
self.assertEqual(len(decompressed), 5)
for i in range(5):
self.assertEqual(decompressed[i].tobytes(), original[i])
def test_single(self):
b = zstd.BufferWithSegments(b'foo', ss.pack(0, 3))
self.assertEqual(len(b), 1)
self.assertEqual(b.size, 3)
self.assertEqual(b.tobytes(), b'foo')
self.assertEqual(len(b[0]), 3)
self.assertEqual(b[0].offset, 0)
self.assertEqual(b[0].tobytes(), b'foo')
def test_invalid_getitem(self):
b = zstd.BufferWithSegments(b'foo', ss.pack(0, 3))
with self.assertRaisesRegexp(IndexError,
'offset must be non-negative'):
test = b[-10]
with self.assertRaisesRegexp(IndexError, 'offset must be less than 1'):
test = b[1]
with self.assertRaisesRegexp(IndexError, 'offset must be less than 1'):
test = b[2]
def test_multiple(self):
b = zstd.BufferWithSegments(
b'foofooxfooxy',
b''.join([ss.pack(0, 3),
ss.pack(3, 4),
ss.pack(7, 5)]))
self.assertEqual(len(b), 3)
self.assertEqual(b.size, 12)
self.assertEqual(b.tobytes(), b'foofooxfooxy')
self.assertEqual(b[0].tobytes(), b'foo')
self.assertEqual(b[1].tobytes(), b'foox')
self.assertEqual(b[2].tobytes(), b'fooxy')
def test_buffer_with_segments_input(self):
cctx = zstd.ZstdCompressor(write_content_size=True,
write_checksum=True)
original = [b'foo' * 4, b'bar' * 6]
frames = [cctx.compress(c) for c in original]
offsets = struct.pack('=QQQQ', 0, len(original[0]), len(original[0]),
len(original[1]))
segments = zstd.BufferWithSegments(b''.join(original), offsets)
result = cctx.multi_compress_to_buffer(segments)
self.assertEqual(len(result), 2)
self.assertEqual(result.size(), 47)
self.assertEqual(result[0].tobytes(), frames[0])
self.assertEqual(result[1].tobytes(), frames[1])
def test_buffer_with_segments_sizes(self):
cctx = zstd.ZstdCompressor(write_content_size=False)
original = [b'foo' * 4, b'bar' * 6, b'baz' * 8]
frames = [cctx.compress(d) for d in original]
sizes = struct.pack('=' + 'Q' * len(original), *map(len, original))
segments = struct.pack('=QQQQQQ', 0, len(frames[0]), len(frames[0]),
len(frames[1]),
len(frames[0]) + len(frames[1]), len(frames[2]))
b = zstd.BufferWithSegments(b''.join(frames), segments)
dctx = zstd.ZstdDecompressor()
result = dctx.multi_decompress_to_buffer(b, decompressed_sizes=sizes)
self.assertEqual(len(result), len(frames))
self.assertEqual(result.size(), sum(map(len, original)))
for i, data in enumerate(original):
self.assertEqual(result[i].tobytes(), data)
def test_buffer_with_segments_input(self):
cctx = zstd.ZstdCompressor(write_content_size=True)
original = [b'foo' * 4, b'bar' * 6]
frames = [cctx.compress(d) for d in original]
dctx = zstd.ZstdDecompressor()
segments = struct.pack('=QQQQ', 0, len(frames[0]), len(frames[0]),
len(frames[1]))
b = zstd.BufferWithSegments(b''.join(frames), segments)
result = dctx.multi_decompress_to_buffer(b)
self.assertEqual(len(result), len(frames))
self.assertEqual(result[0].offset, 0)
self.assertEqual(len(result[0]), 12)
self.assertEqual(result[1].offset, 12)
self.assertEqual(len(result[1]), 18)
def test_invalid_offset(self):
with self.assertRaisesRegexp(
ValueError, 'offset within segments array references memory'):
zstd.BufferWithSegments(b'foo', ss.pack(0, 4))