Google Git Google Git
Sign in
chromium/external/github.com/python/cpython/626394e0a38bdf7734567f00a09837f09e653bff/./Lib/test/test_zstd.py
blob: 372326a26c71eaf89146f832d9a20c33b46c689d [file]
import array
import gc
import io
import pathlib
import random
import re
import os
import unittest
import tempfile
import threading
from test import support
from test.support.import_helper import import_module
from test.support import threading_helper
from test.support import _1M
_zstd = import_module("_zstd")
zstd = import_module("compression.zstd")
from compression.zstd import (
open,
compress,
decompress,
ZstdCompressor,
ZstdDecompressor,
ZstdDict,
ZstdError,
COMPRESSION_LEVEL_DEFAULT,
get_frame_info,
get_frame_size,
finalize_dict,
train_dict,
CompressionParameter,
DecompressionParameter,
Strategy,
ZstdFile,
)
_1K = 1024
_130_1K = 130 * _1K
DICT_SIZE1 = 3*_1K
DAT_130K_D = None
DAT_130K_C = None
DECOMPRESSED_DAT = None
COMPRESSED_DAT = None
DECOMPRESSED_100_PLUS_32KB = None
COMPRESSED_100_PLUS_32KB = None
SKIPPABLE_FRAME = None
THIS_FILE_BYTES = None
THIS_FILE_STR = None
COMPRESSED_THIS_FILE = None
COMPRESSED_BOGUS = None
SAMPLES = None
TRAINED_DICT = None
# Cannot be deferred to setup as it is used to check whether or not to skip
# tests
try:
SUPPORT_MULTITHREADING = CompressionParameter.nb_workers.bounds() != (0, 0)
except Exception:
SUPPORT_MULTITHREADING = False
C_INT_MIN = -(2**31)
C_INT_MAX = (2**31) - 1
def setUpModule():
# uncompressed size 130KB, more than a zstd block.
# with a frame epilogue, 4 bytes checksum.
global DAT_130K_D
DAT_130K_D = bytes([random.randint(0, 127) for _ in range(130*_1K)])
global DAT_130K_C
DAT_130K_C = compress(DAT_130K_D, options={CompressionParameter.checksum_flag:1})
global DECOMPRESSED_DAT
DECOMPRESSED_DAT = b'abcdefg123456' * 1000
global COMPRESSED_DAT
COMPRESSED_DAT = compress(DECOMPRESSED_DAT)
global DECOMPRESSED_100_PLUS_32KB
DECOMPRESSED_100_PLUS_32KB = b'a' * (100 + 32*_1K)
global COMPRESSED_100_PLUS_32KB
COMPRESSED_100_PLUS_32KB = compress(DECOMPRESSED_100_PLUS_32KB)
global SKIPPABLE_FRAME
SKIPPABLE_FRAME = (0x184D2A50).to_bytes(4, byteorder='little') + \
(32*_1K).to_bytes(4, byteorder='little') + \
b'a' * (32*_1K)
global THIS_FILE_BYTES, THIS_FILE_STR
with io.open(os.path.abspath(__file__), 'rb') as f:
THIS_FILE_BYTES = f.read()
THIS_FILE_BYTES = re.sub(rb'\r?\n', rb'\n', THIS_FILE_BYTES)
THIS_FILE_STR = THIS_FILE_BYTES.decode('utf-8')
global COMPRESSED_THIS_FILE
COMPRESSED_THIS_FILE = compress(THIS_FILE_BYTES)
global COMPRESSED_BOGUS
COMPRESSED_BOGUS = DECOMPRESSED_DAT
# dict data
words = [b'red', b'green', b'yellow', b'black', b'withe', b'blue',
b'lilac', b'purple', b'navy', b'glod', b'silver', b'olive',
b'dog', b'cat', b'tiger', b'lion', b'fish', b'bird']
lst = []
for i in range(300):
sample = [b'%s = %d' % (random.choice(words), random.randrange(100))
for j in range(20)]
sample = b'\n'.join(sample)
lst.append(sample)
global SAMPLES
SAMPLES = lst
assert len(SAMPLES) > 10
global TRAINED_DICT
TRAINED_DICT = train_dict(SAMPLES, 3*_1K)
assert len(TRAINED_DICT.dict_content) <= 3*_1K
class FunctionsTestCase(unittest.TestCase):
def _test_zstd_version(self, v, string):
self.assertIsInstance(v[:], tuple)
self.assertEqual(len(v), 3)
self.assertIsInstance(v[0], int)
self.assertIsInstance(v[1], int)
self.assertIsInstance(v[2], int)
self.assertIsInstance(v.major, int)
self.assertIsInstance(v.minor, int)
self.assertIsInstance(v.patch, int)
self.assertEqual(v[0], v.major)
self.assertEqual(v[1], v.minor)
self.assertEqual(v[2], v.patch)
self.assertGreaterEqual(v.major, 0)
self.assertGreaterEqual(v.minor, 0)
self.assertGreaterEqual(v.patch, 0)
self.assertEqual(string, '%d.%d.%d' % v)
def test_zstd_version(self):
if support.verbose:
print(f'ZSTD_VERSION = {zstd.ZSTD_VERSION}', flush=True)
print(f'zstd_version = {zstd.zstd_version}', flush=True)
print(f'ZSTD_VERSION_INFO = {zstd.ZSTD_VERSION_INFO}', flush=True)
print(f'zstd_version_info = {zstd.zstd_version_info}', flush=True)
self._test_zstd_version(zstd.ZSTD_VERSION_INFO, zstd.ZSTD_VERSION)
self._test_zstd_version(zstd.zstd_version_info, zstd.zstd_version)
self.assertEqual(zstd.ZSTD_VERSION_INFO[0], zstd.zstd_version_info[0])
def test_compressionLevel_values(self):
min, max = CompressionParameter.compression_level.bounds()
self.assertIs(type(COMPRESSION_LEVEL_DEFAULT), int)
self.assertIs(type(min), int)
self.assertIs(type(max), int)
self.assertLess(min, max)
def test_roundtrip_default(self):
raw_dat = THIS_FILE_BYTES[: len(THIS_FILE_BYTES) // 6]
dat1 = compress(raw_dat)
dat2 = decompress(dat1)
self.assertEqual(dat2, raw_dat)
def test_roundtrip_level(self):
raw_dat = THIS_FILE_BYTES[: len(THIS_FILE_BYTES) // 6]
level_min, level_max = CompressionParameter.compression_level.bounds()
for level in range(max(-20, level_min), level_max + 1):
dat1 = compress(raw_dat, level)
dat2 = decompress(dat1)
self.assertEqual(dat2, raw_dat)
def test_get_frame_info(self):
# no dict
info = get_frame_info(COMPRESSED_100_PLUS_32KB[:20])
self.assertEqual(info.decompressed_size, 32 * _1K + 100)
self.assertEqual(info.dictionary_id, 0)
# use dict
dat = compress(b"a" * 345, zstd_dict=TRAINED_DICT)
info = get_frame_info(dat)
self.assertEqual(info.decompressed_size, 345)
self.assertEqual(info.dictionary_id, TRAINED_DICT.dict_id)
with self.assertRaisesRegex(ZstdError, "not less than the frame header"):
get_frame_info(b"aaaaaaaaaaaaaa")
def test_get_frame_size(self):
size = get_frame_size(COMPRESSED_100_PLUS_32KB)
self.assertEqual(size, len(COMPRESSED_100_PLUS_32KB))
with self.assertRaisesRegex(ZstdError, "not less than this complete frame"):
get_frame_size(b"aaaaaaaaaaaaaa")
def test_decompress_2x130_1K(self):
decompressed_size = get_frame_info(DAT_130K_C).decompressed_size
self.assertEqual(decompressed_size, _130_1K)
dat = decompress(DAT_130K_C + DAT_130K_C)
self.assertEqual(len(dat), 2 * _130_1K)
class CompressorTestCase(unittest.TestCase):
def test_simple_compress_bad_args(self):
# ZstdCompressor
self.assertRaises(TypeError, ZstdCompressor, [])
self.assertRaises(TypeError, ZstdCompressor, level=3.14)
self.assertRaises(TypeError, ZstdCompressor, level="abc")
self.assertRaises(TypeError, ZstdCompressor, options=b"abc")
self.assertRaises(TypeError, ZstdCompressor, zstd_dict=123)
self.assertRaises(TypeError, ZstdCompressor, zstd_dict=b"abcd1234")
self.assertRaises(TypeError, ZstdCompressor, zstd_dict={1: 2, 3: 4})
# valid range for compression level is [-(1<<17), 22]
msg = r'illegal compression level {}; the valid range is \[-?\d+, -?\d+\]'
with self.assertRaisesRegex(ValueError, msg.format(C_INT_MAX)):
ZstdCompressor(C_INT_MAX)
with self.assertRaisesRegex(ValueError, msg.format(C_INT_MIN)):
ZstdCompressor(C_INT_MIN)
msg = r'illegal compression level; the valid range is \[-?\d+, -?\d+\]'
with self.assertRaisesRegex(ValueError, msg):
ZstdCompressor(level=-(2**1000))
with self.assertRaisesRegex(ValueError, msg):
ZstdCompressor(level=2**1000)
with self.assertRaises(ValueError):
ZstdCompressor(options={CompressionParameter.window_log: 100})
with self.assertRaises(ValueError):
ZstdCompressor(options={3333: 100})
# Method bad arguments
zc = ZstdCompressor()
self.assertRaises(TypeError, zc.compress)
self.assertRaises((TypeError, ValueError), zc.compress, b"foo", b"bar")
self.assertRaises(TypeError, zc.compress, "str")
self.assertRaises((TypeError, ValueError), zc.flush, b"foo")
self.assertRaises(TypeError, zc.flush, b"blah", 1)
self.assertRaises(ValueError, zc.compress, b'', -1)
self.assertRaises(ValueError, zc.compress, b'', 3)
self.assertRaises(ValueError, zc.flush, zc.CONTINUE) # 0
self.assertRaises(ValueError, zc.flush, 3)
zc.compress(b'')
zc.compress(b'', zc.CONTINUE)
zc.compress(b'', zc.FLUSH_BLOCK)
zc.compress(b'', zc.FLUSH_FRAME)
empty = zc.flush()
zc.flush(zc.FLUSH_BLOCK)
zc.flush(zc.FLUSH_FRAME)
def test_compress_parameters(self):
d = {CompressionParameter.compression_level : 10,
CompressionParameter.window_log : 12,
CompressionParameter.hash_log : 10,
CompressionParameter.chain_log : 12,
CompressionParameter.search_log : 12,
CompressionParameter.min_match : 4,
CompressionParameter.target_length : 12,
CompressionParameter.strategy : Strategy.lazy,
CompressionParameter.enable_long_distance_matching : 1,
CompressionParameter.ldm_hash_log : 12,
CompressionParameter.ldm_min_match : 11,
CompressionParameter.ldm_bucket_size_log : 5,
CompressionParameter.ldm_hash_rate_log : 12,
CompressionParameter.content_size_flag : 1,
CompressionParameter.checksum_flag : 1,
CompressionParameter.dict_id_flag : 0,
CompressionParameter.nb_workers : 2 if SUPPORT_MULTITHREADING else 0,
CompressionParameter.job_size : 5*_1M if SUPPORT_MULTITHREADING else 0,
CompressionParameter.overlap_log : 9 if SUPPORT_MULTITHREADING else 0,
}
ZstdCompressor(options=d)
d1 = d.copy()
# larger than signed int
d1[CompressionParameter.ldm_bucket_size_log] = C_INT_MAX
with self.assertRaises(ValueError):
ZstdCompressor(options=d1)
# smaller than signed int
d1[CompressionParameter.ldm_bucket_size_log] = C_INT_MIN
with self.assertRaises(ValueError):
ZstdCompressor(options=d1)
# out of bounds compression level
level_min, level_max = CompressionParameter.compression_level.bounds()
with self.assertRaises(ValueError):
compress(b'', level_max+1)
with self.assertRaises(ValueError):
compress(b'', level_min-1)
with self.assertRaises(ValueError):
compress(b'', 2**1000)
with self.assertRaises(ValueError):
compress(b'', -(2**1000))
with self.assertRaises(ValueError):
compress(b'', options={
CompressionParameter.compression_level: level_max+1})
with self.assertRaises(ValueError):
compress(b'', options={
CompressionParameter.compression_level: level_min-1})
# zstd lib doesn't support MT compression
if not SUPPORT_MULTITHREADING:
with self.assertRaises(ValueError):
ZstdCompressor(options={CompressionParameter.nb_workers:4})
with self.assertRaises(ValueError):
ZstdCompressor(options={CompressionParameter.job_size:4})
with self.assertRaises(ValueError):
ZstdCompressor(options={CompressionParameter.overlap_log:4})
# out of bounds error msg
option = {CompressionParameter.window_log:100}
with self.assertRaisesRegex(
ValueError,
"compression parameter 'window_log' received an illegal value 100; "
r'the valid range is \[-?\d+, -?\d+\]',
):
compress(b'', options=option)
def test_unknown_compression_parameter(self):
KEY = 100001234
option = {CompressionParameter.compression_level: 10,
KEY: 200000000}
pattern = rf"invalid compression parameter 'unknown parameter \(key {KEY}\)'"
with self.assertRaisesRegex(ValueError, pattern):
ZstdCompressor(options=option)
@unittest.skipIf(not SUPPORT_MULTITHREADING,
"zstd build doesn't support multi-threaded compression")
def test_zstd_multithread_compress(self):
size = 40*_1M
b = THIS_FILE_BYTES * (size // len(THIS_FILE_BYTES))
options = {CompressionParameter.compression_level : 4,
CompressionParameter.nb_workers : 2}
# compress()
dat1 = compress(b, options=options)
dat2 = decompress(dat1)
self.assertEqual(dat2, b)
# ZstdCompressor
c = ZstdCompressor(options=options)
dat1 = c.compress(b, c.CONTINUE)
dat2 = c.compress(b, c.FLUSH_BLOCK)
dat3 = c.compress(b, c.FLUSH_FRAME)
dat4 = decompress(dat1+dat2+dat3)
self.assertEqual(dat4, b * 3)
# ZstdFile
with ZstdFile(io.BytesIO(), 'w', options=options) as f:
f.write(b)
def test_compress_flushblock(self):
point = len(THIS_FILE_BYTES) // 2
c = ZstdCompressor()
self.assertEqual(c.last_mode, c.FLUSH_FRAME)
dat1 = c.compress(THIS_FILE_BYTES[:point])
self.assertEqual(c.last_mode, c.CONTINUE)
dat1 += c.compress(THIS_FILE_BYTES[point:], c.FLUSH_BLOCK)
self.assertEqual(c.last_mode, c.FLUSH_BLOCK)
dat2 = c.flush()
pattern = "Compressed data ended before the end-of-stream marker"
with self.assertRaisesRegex(ZstdError, pattern):
decompress(dat1)
dat3 = decompress(dat1 + dat2)
self.assertEqual(dat3, THIS_FILE_BYTES)
def test_compress_flushframe(self):
# test compress & decompress
point = len(THIS_FILE_BYTES) // 2
c = ZstdCompressor()
dat1 = c.compress(THIS_FILE_BYTES[:point])
self.assertEqual(c.last_mode, c.CONTINUE)
dat1 += c.compress(THIS_FILE_BYTES[point:], c.FLUSH_FRAME)
self.assertEqual(c.last_mode, c.FLUSH_FRAME)
nt = get_frame_info(dat1)
self.assertEqual(nt.decompressed_size, None) # no content size
dat2 = decompress(dat1)
self.assertEqual(dat2, THIS_FILE_BYTES)
# single .FLUSH_FRAME mode has content size
c = ZstdCompressor()
dat = c.compress(THIS_FILE_BYTES, mode=c.FLUSH_FRAME)
self.assertEqual(c.last_mode, c.FLUSH_FRAME)
nt = get_frame_info(dat)
self.assertEqual(nt.decompressed_size, len(THIS_FILE_BYTES))
def test_compress_empty(self):
# output empty content frame
self.assertNotEqual(compress(b''), b'')
c = ZstdCompressor()
self.assertNotEqual(c.compress(b'', c.FLUSH_FRAME), b'')
def test_set_pledged_input_size(self):
DAT = DECOMPRESSED_100_PLUS_32KB
CHUNK_SIZE = len(DAT) // 3
# wrong value
c = ZstdCompressor()
with self.assertRaisesRegex(ValueError,
r'should be a positive int less than \d+'):
c.set_pledged_input_size(-300)
# overflow
with self.assertRaisesRegex(ValueError,
r'should be a positive int less than \d+'):
c.set_pledged_input_size(2**64)
# ZSTD_CONTENTSIZE_ERROR is invalid
with self.assertRaisesRegex(ValueError,
r'should be a positive int less than \d+'):
c.set_pledged_input_size(2**64-2)
# ZSTD_CONTENTSIZE_UNKNOWN should use None
with self.assertRaisesRegex(ValueError,
r'should be a positive int less than \d+'):
c.set_pledged_input_size(2**64-1)
# check valid values are settable
c.set_pledged_input_size(2**63)
c.set_pledged_input_size(2**64-3)
# check that zero means empty frame
c = ZstdCompressor(level=1)
c.set_pledged_input_size(0)
c.compress(b'')
dat = c.flush()
ret = get_frame_info(dat)
self.assertEqual(ret.decompressed_size, 0)
# wrong mode
c = ZstdCompressor(level=1)
c.compress(b'123456')
self.assertEqual(c.last_mode, c.CONTINUE)
with self.assertRaisesRegex(ValueError,
r'last_mode == FLUSH_FRAME'):
c.set_pledged_input_size(300)
# None value
c = ZstdCompressor(level=1)
c.set_pledged_input_size(None)
dat = c.compress(DAT) + c.flush()
ret = get_frame_info(dat)
self.assertEqual(ret.decompressed_size, None)
# correct value
c = ZstdCompressor(level=1)
c.set_pledged_input_size(len(DAT))
chunks = []
posi = 0
while posi < len(DAT):
dat = c.compress(DAT[posi:posi+CHUNK_SIZE])
posi += CHUNK_SIZE
chunks.append(dat)
dat = c.flush()
chunks.append(dat)
chunks = b''.join(chunks)
ret = get_frame_info(chunks)
self.assertEqual(ret.decompressed_size, len(DAT))
self.assertEqual(decompress(chunks), DAT)
c.set_pledged_input_size(len(DAT)) # the second frame
dat = c.compress(DAT) + c.flush()
ret = get_frame_info(dat)
self.assertEqual(ret.decompressed_size, len(DAT))
self.assertEqual(decompress(dat), DAT)
# not enough data
c = ZstdCompressor(level=1)
c.set_pledged_input_size(len(DAT)+1)
for start in range(0, len(DAT), CHUNK_SIZE):
end = min(start+CHUNK_SIZE, len(DAT))
_dat = c.compress(DAT[start:end])
with self.assertRaises(ZstdError):
c.flush()
# too much data
c = ZstdCompressor(level=1)
c.set_pledged_input_size(len(DAT))
for start in range(0, len(DAT), CHUNK_SIZE):
end = min(start+CHUNK_SIZE, len(DAT))
_dat = c.compress(DAT[start:end])
with self.assertRaises(ZstdError):
c.compress(b'extra', ZstdCompressor.FLUSH_FRAME)
# content size not set if content_size_flag == 0
c = ZstdCompressor(options={CompressionParameter.content_size_flag: 0})
c.set_pledged_input_size(10)
dat1 = c.compress(b"hello")
dat2 = c.compress(b"world")
dat3 = c.flush()
frame_data = get_frame_info(dat1 + dat2 + dat3)
self.assertIsNone(frame_data.decompressed_size)
class DecompressorTestCase(unittest.TestCase):
def test_simple_decompress_bad_args(self):
# ZstdDecompressor
self.assertRaises(TypeError, ZstdDecompressor, ())
self.assertRaises(TypeError, ZstdDecompressor, zstd_dict=123)
self.assertRaises(TypeError, ZstdDecompressor, zstd_dict=b'abc')
self.assertRaises(TypeError, ZstdDecompressor, zstd_dict={1:2, 3:4})
self.assertRaises(TypeError, ZstdDecompressor, options=123)
self.assertRaises(TypeError, ZstdDecompressor, options='abc')
self.assertRaises(TypeError, ZstdDecompressor, options=b'abc')
with self.assertRaises(ValueError):
ZstdDecompressor(options={C_INT_MAX: 100})
with self.assertRaises(ValueError):
ZstdDecompressor(options={C_INT_MIN: 100})
with self.assertRaises(ValueError):
ZstdDecompressor(options={0: C_INT_MAX})
with self.assertRaises(OverflowError):
ZstdDecompressor(options={2**1000: 100})
with self.assertRaises(OverflowError):
ZstdDecompressor(options={-(2**1000): 100})
with self.assertRaises(OverflowError):
ZstdDecompressor(options={0: -(2**1000)})
with self.assertRaises(ValueError):
ZstdDecompressor(options={DecompressionParameter.window_log_max: 100})
with self.assertRaises(ValueError):
ZstdDecompressor(options={3333: 100})
empty = compress(b'')
lzd = ZstdDecompressor()
self.assertRaises(TypeError, lzd.decompress)
self.assertRaises(TypeError, lzd.decompress, b"foo", b"bar")
self.assertRaises(TypeError, lzd.decompress, "str")
lzd.decompress(empty)
def test_decompress_parameters(self):
d = {DecompressionParameter.window_log_max : 15}
ZstdDecompressor(options=d)
d1 = d.copy()
# larger than signed int
d1[DecompressionParameter.window_log_max] = 2**1000
with self.assertRaises(OverflowError):
ZstdDecompressor(None, d1)
# smaller than signed int
d1[DecompressionParameter.window_log_max] = -(2**1000)
with self.assertRaises(OverflowError):
ZstdDecompressor(None, d1)
d1[DecompressionParameter.window_log_max] = C_INT_MAX
with self.assertRaises(ValueError):
ZstdDecompressor(None, d1)
d1[DecompressionParameter.window_log_max] = C_INT_MIN
with self.assertRaises(ValueError):
ZstdDecompressor(None, d1)
# out of bounds error msg
options = {DecompressionParameter.window_log_max:100}
with self.assertRaisesRegex(
ValueError,
"decompression parameter 'window_log_max' received an illegal value 100; "
r'the valid range is \[-?\d+, -?\d+\]',
):
decompress(b'', options=options)
# out of bounds deecompression parameter
options[DecompressionParameter.window_log_max] = C_INT_MAX
with self.assertRaises(ValueError):
decompress(b'', options=options)
options[DecompressionParameter.window_log_max] = C_INT_MIN
with self.assertRaises(ValueError):
decompress(b'', options=options)
options[DecompressionParameter.window_log_max] = 2**1000
with self.assertRaises(OverflowError):
decompress(b'', options=options)
options[DecompressionParameter.window_log_max] = -(2**1000)
with self.assertRaises(OverflowError):
decompress(b'', options=options)
def test_unknown_decompression_parameter(self):
KEY = 100001234
options = {DecompressionParameter.window_log_max: DecompressionParameter.window_log_max.bounds()[1],
KEY: 200000000}
pattern = rf"invalid decompression parameter 'unknown parameter \(key {KEY}\)'"
with self.assertRaisesRegex(ValueError, pattern):
ZstdDecompressor(options=options)
def test_decompress_epilogue_flags(self):
# DAT_130K_C has a 4 bytes checksum at frame epilogue
# full unlimited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.needs_input)
with self.assertRaises(EOFError):
dat = d.decompress(b'')
# full limited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C, _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.needs_input)
with self.assertRaises(EOFError):
dat = d.decompress(b'', 0)
# [:-4] unlimited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-4])
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.needs_input)
dat = d.decompress(b'')
self.assertEqual(len(dat), 0)
self.assertTrue(d.needs_input)
# [:-4] limited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-4], _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.needs_input)
dat = d.decompress(b'', 0)
self.assertEqual(len(dat), 0)
self.assertFalse(d.needs_input)
# [:-3] unlimited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-3])
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.needs_input)
dat = d.decompress(b'')
self.assertEqual(len(dat), 0)
self.assertTrue(d.needs_input)
# [:-3] limited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-3], _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.needs_input)
dat = d.decompress(b'', 0)
self.assertEqual(len(dat), 0)
self.assertFalse(d.needs_input)
# [:-1] unlimited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-1])
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.needs_input)
dat = d.decompress(b'')
self.assertEqual(len(dat), 0)
self.assertTrue(d.needs_input)
# [:-1] limited
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-1], _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.needs_input)
dat = d.decompress(b'', 0)
self.assertEqual(len(dat), 0)
self.assertFalse(d.needs_input)
def test_decompressor_arg(self):
zd = ZstdDict(b'12345678', is_raw=True)
with self.assertRaises(TypeError):
d = ZstdDecompressor(zstd_dict={})
with self.assertRaises(TypeError):
d = ZstdDecompressor(options=zd)
ZstdDecompressor()
ZstdDecompressor(zd, {})
ZstdDecompressor(zstd_dict=zd, options={DecompressionParameter.window_log_max:25})
def test_decompressor_1(self):
# empty
d = ZstdDecompressor()
dat = d.decompress(b'')
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
# 130_1K full
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C)
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
# 130_1K full, limit output
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C, _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
# 130_1K, without 4 bytes checksum
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-4])
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
# above, limit output
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C[:-4], _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
# full, unused_data
TRAIL = b'89234893abcd'
d = ZstdDecompressor()
dat = d.decompress(DAT_130K_C + TRAIL, _130_1K)
self.assertEqual(len(dat), _130_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, TRAIL)
def test_decompressor_chunks_read_300(self):
TRAIL = b'89234893abcd'
DAT = DAT_130K_C + TRAIL
d = ZstdDecompressor()
bi = io.BytesIO(DAT)
lst = []
while True:
if d.needs_input:
dat = bi.read(300)
if not dat:
break
else:
raise Exception('should not get here')
ret = d.decompress(dat)
lst.append(ret)
if d.eof:
break
ret = b''.join(lst)
self.assertEqual(len(ret), _130_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data + bi.read(), TRAIL)
def test_decompressor_chunks_read_3(self):
TRAIL = b'89234893'
DAT = DAT_130K_C + TRAIL
d = ZstdDecompressor()
bi = io.BytesIO(DAT)
lst = []
while True:
if d.needs_input:
dat = bi.read(3)
if not dat:
break
else:
dat = b''
ret = d.decompress(dat, 1)
lst.append(ret)
if d.eof:
break
ret = b''.join(lst)
self.assertEqual(len(ret), _130_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data + bi.read(), TRAIL)
def test_decompress_empty(self):
with self.assertRaises(ZstdError):
decompress(b'')
d = ZstdDecompressor()
self.assertEqual(d.decompress(b''), b'')
self.assertFalse(d.eof)
def test_decompress_empty_content_frame(self):
DAT = compress(b'')
# decompress
self.assertGreaterEqual(len(DAT), 4)
self.assertEqual(decompress(DAT), b'')
with self.assertRaises(ZstdError):
decompress(DAT[:-1])
# ZstdDecompressor
d = ZstdDecompressor()
dat = d.decompress(DAT)
self.assertEqual(dat, b'')
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
d = ZstdDecompressor()
dat = d.decompress(DAT[:-1])
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
class DecompressorFlagsTestCase(unittest.TestCase):
@classmethod
def setUpClass(cls):
options = {CompressionParameter.checksum_flag:1}
c = ZstdCompressor(options=options)
cls.DECOMPRESSED_42 = b'a'*42
cls.FRAME_42 = c.compress(cls.DECOMPRESSED_42, c.FLUSH_FRAME)
cls.DECOMPRESSED_60 = b'a'*60
cls.FRAME_60 = c.compress(cls.DECOMPRESSED_60, c.FLUSH_FRAME)
cls.FRAME_42_60 = cls.FRAME_42 + cls.FRAME_60
cls.DECOMPRESSED_42_60 = cls.DECOMPRESSED_42 + cls.DECOMPRESSED_60
cls._130_1K = 130*_1K
c = ZstdCompressor()
cls.UNKNOWN_FRAME_42 = c.compress(cls.DECOMPRESSED_42) + c.flush()
cls.UNKNOWN_FRAME_60 = c.compress(cls.DECOMPRESSED_60) + c.flush()
cls.UNKNOWN_FRAME_42_60 = cls.UNKNOWN_FRAME_42 + cls.UNKNOWN_FRAME_60
cls.TRAIL = b'12345678abcdefg!@#$%^&*()_+|'
def test_function_decompress(self):
self.assertEqual(len(decompress(COMPRESSED_100_PLUS_32KB)), 100+32*_1K)
# 1 frame
self.assertEqual(decompress(self.FRAME_42), self.DECOMPRESSED_42)
self.assertEqual(decompress(self.UNKNOWN_FRAME_42), self.DECOMPRESSED_42)
pattern = r"Compressed data ended before the end-of-stream marker"
with self.assertRaisesRegex(ZstdError, pattern):
decompress(self.FRAME_42[:1])
with self.assertRaisesRegex(ZstdError, pattern):
decompress(self.FRAME_42[:-4])
with self.assertRaisesRegex(ZstdError, pattern):
decompress(self.FRAME_42[:-1])
# 2 frames
self.assertEqual(decompress(self.FRAME_42_60), self.DECOMPRESSED_42_60)
self.assertEqual(decompress(self.UNKNOWN_FRAME_42_60), self.DECOMPRESSED_42_60)
self.assertEqual(decompress(self.FRAME_42 + self.UNKNOWN_FRAME_60),
self.DECOMPRESSED_42_60)
self.assertEqual(decompress(self.UNKNOWN_FRAME_42 + self.FRAME_60),
self.DECOMPRESSED_42_60)
with self.assertRaisesRegex(ZstdError, pattern):
decompress(self.FRAME_42_60[:-4])
with self.assertRaisesRegex(ZstdError, pattern):
decompress(self.UNKNOWN_FRAME_42_60[:-1])
# 130_1K
self.assertEqual(decompress(DAT_130K_C), DAT_130K_D)
with self.assertRaisesRegex(ZstdError, pattern):
decompress(DAT_130K_C[:-4])
with self.assertRaisesRegex(ZstdError, pattern):
decompress(DAT_130K_C[:-1])
# Unknown frame descriptor
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(b'aaaaaaaaa')
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(self.FRAME_42 + b'aaaaaaaaa')
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(self.UNKNOWN_FRAME_42_60 + b'aaaaaaaaa')
# doesn't match checksum
checksum = DAT_130K_C[-4:]
if checksum[0] == 255:
wrong_checksum = bytes([254]) + checksum[1:]
else:
wrong_checksum = bytes([checksum[0]+1]) + checksum[1:]
dat = DAT_130K_C[:-4] + wrong_checksum
with self.assertRaisesRegex(ZstdError, "doesn't match checksum"):
decompress(dat)
def test_function_skippable(self):
self.assertEqual(decompress(SKIPPABLE_FRAME), b'')
self.assertEqual(decompress(SKIPPABLE_FRAME + SKIPPABLE_FRAME), b'')
# 1 frame + 2 skippable
self.assertEqual(len(decompress(SKIPPABLE_FRAME + SKIPPABLE_FRAME + DAT_130K_C)),
self._130_1K)
self.assertEqual(len(decompress(DAT_130K_C + SKIPPABLE_FRAME + SKIPPABLE_FRAME)),
self._130_1K)
self.assertEqual(len(decompress(SKIPPABLE_FRAME + DAT_130K_C + SKIPPABLE_FRAME)),
self._130_1K)
# unknown size
self.assertEqual(decompress(SKIPPABLE_FRAME + self.UNKNOWN_FRAME_60),
self.DECOMPRESSED_60)
self.assertEqual(decompress(self.UNKNOWN_FRAME_60 + SKIPPABLE_FRAME),
self.DECOMPRESSED_60)
# 2 frames + 1 skippable
self.assertEqual(decompress(self.FRAME_42 + SKIPPABLE_FRAME + self.FRAME_60),
self.DECOMPRESSED_42_60)
self.assertEqual(decompress(SKIPPABLE_FRAME + self.FRAME_42_60),
self.DECOMPRESSED_42_60)
self.assertEqual(decompress(self.UNKNOWN_FRAME_42_60 + SKIPPABLE_FRAME),
self.DECOMPRESSED_42_60)
# incomplete
with self.assertRaises(ZstdError):
decompress(SKIPPABLE_FRAME[:1])
with self.assertRaises(ZstdError):
decompress(SKIPPABLE_FRAME[:-1])
with self.assertRaises(ZstdError):
decompress(self.FRAME_42 + SKIPPABLE_FRAME[:-1])
# Unknown frame descriptor
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(b'aaaaaaaaa' + SKIPPABLE_FRAME)
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(SKIPPABLE_FRAME + b'aaaaaaaaa')
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
decompress(SKIPPABLE_FRAME + SKIPPABLE_FRAME + b'aaaaaaaaa')
def test_decompressor_1(self):
# empty 1
d = ZstdDecompressor()
dat = d.decompress(b'')
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(b'', 0)
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(COMPRESSED_100_PLUS_32KB + b'a')
self.assertEqual(dat, DECOMPRESSED_100_PLUS_32KB)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'a')
self.assertEqual(d.unused_data, b'a') # twice
# empty 2
d = ZstdDecompressor()
dat = d.decompress(b'', 0)
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(b'')
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(COMPRESSED_100_PLUS_32KB + b'a')
self.assertEqual(dat, DECOMPRESSED_100_PLUS_32KB)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'a')
self.assertEqual(d.unused_data, b'a') # twice
# 1 frame
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_42)
self.assertEqual(dat, self.DECOMPRESSED_42)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
with self.assertRaises(EOFError):
d.decompress(b'')
# 1 frame, trail
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_42 + self.TRAIL)
self.assertEqual(dat, self.DECOMPRESSED_42)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, self.TRAIL)
self.assertEqual(d.unused_data, self.TRAIL) # twice
# 1 frame, 32_1K
temp = compress(b'a'*(32*_1K))
d = ZstdDecompressor()
dat = d.decompress(temp, 32*_1K)
self.assertEqual(dat, b'a'*(32*_1K))
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
with self.assertRaises(EOFError):
d.decompress(b'')
# 1 frame, 32_1K+100, trail
d = ZstdDecompressor()
dat = d.decompress(COMPRESSED_100_PLUS_32KB+self.TRAIL, 100) # 100 bytes
self.assertEqual(len(dat), 100)
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
dat = d.decompress(b'') # 32_1K
self.assertEqual(len(dat), 32*_1K)
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, self.TRAIL)
self.assertEqual(d.unused_data, self.TRAIL) # twice
with self.assertRaises(EOFError):
d.decompress(b'')
# incomplete 1
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_60[:1])
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# incomplete 2
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_60[:-4])
self.assertEqual(dat, self.DECOMPRESSED_60)
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# incomplete 3
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_60[:-1])
self.assertEqual(dat, self.DECOMPRESSED_60)
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
# incomplete 4
d = ZstdDecompressor()
dat = d.decompress(self.FRAME_60[:-4], 60)
self.assertEqual(dat, self.DECOMPRESSED_60)
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(b'')
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# Unknown frame descriptor
d = ZstdDecompressor()
with self.assertRaisesRegex(ZstdError, "Unknown frame descriptor"):
d.decompress(b'aaaaaaaaa')
def test_decompressor_skippable(self):
# 1 skippable
d = ZstdDecompressor()
dat = d.decompress(SKIPPABLE_FRAME)
self.assertEqual(dat, b'')
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# 1 skippable, max_length=0
d = ZstdDecompressor()
dat = d.decompress(SKIPPABLE_FRAME, 0)
self.assertEqual(dat, b'')
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# 1 skippable, trail
d = ZstdDecompressor()
dat = d.decompress(SKIPPABLE_FRAME + self.TRAIL)
self.assertEqual(dat, b'')
self.assertTrue(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, self.TRAIL)
self.assertEqual(d.unused_data, self.TRAIL) # twice
# incomplete
d = ZstdDecompressor()
dat = d.decompress(SKIPPABLE_FRAME[:-1])
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
# incomplete
d = ZstdDecompressor()
dat = d.decompress(SKIPPABLE_FRAME[:-1], 0)
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertFalse(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
dat = d.decompress(b'')
self.assertEqual(dat, b'')
self.assertFalse(d.eof)
self.assertTrue(d.needs_input)
self.assertEqual(d.unused_data, b'')
self.assertEqual(d.unused_data, b'') # twice
class ZstdDictTestCase(unittest.TestCase):
def test_is_raw(self):
# must be passed as a keyword argument
with self.assertRaises(TypeError):
ZstdDict(bytes(8), True)
# content < 8
b = b'1234567'
with self.assertRaises(ValueError):
ZstdDict(b)
# content == 8
b = b'12345678'
zd = ZstdDict(b, is_raw=True)
self.assertEqual(zd.dict_id, 0)
temp = compress(b'aaa12345678', level=3, zstd_dict=zd)
self.assertEqual(b'aaa12345678', decompress(temp, zd))
# is_raw == False
b = b'12345678abcd'
with self.assertRaises(ValueError):
ZstdDict(b)
# read only attributes
with self.assertRaises(AttributeError):
zd.dict_content = b
with self.assertRaises(AttributeError):
zd.dict_id = 10000
# ZstdDict arguments
zd = ZstdDict(TRAINED_DICT.dict_content, is_raw=False)
self.assertNotEqual(zd.dict_id, 0)
zd = ZstdDict(TRAINED_DICT.dict_content, is_raw=True)
self.assertNotEqual(zd.dict_id, 0) # note this assertion
with self.assertRaises(TypeError):
ZstdDict("12345678abcdef", is_raw=True)
with self.assertRaises(TypeError):
ZstdDict(TRAINED_DICT)
# invalid parameter
with self.assertRaises(TypeError):
ZstdDict(desk333=345)
def test_invalid_dict(self):
DICT_MAGIC = 0xEC30A437.to_bytes(4, byteorder='little')
dict_content = DICT_MAGIC + b'abcdefghighlmnopqrstuvwxyz'
# corrupted
zd = ZstdDict(dict_content, is_raw=False)
with self.assertRaisesRegex(ZstdError, r'ZSTD_CDict.*?content\.$'):
ZstdCompressor(zstd_dict=zd.as_digested_dict)
with self.assertRaisesRegex(ZstdError, r'ZSTD_DDict.*?content\.$'):
ZstdDecompressor(zd)
# wrong type
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=[zd, 1])
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=(zd, 1.0))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=(zd,))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=(zd, 1, 2))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=(zd, -1))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdCompressor(zstd_dict=(zd, 3))
with self.assertRaises(OverflowError):
ZstdCompressor(zstd_dict=(zd, 2**1000))
with self.assertRaises(OverflowError):
ZstdCompressor(zstd_dict=(zd, -2**1000))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor(zstd_dict=[zd, 1])
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor(zstd_dict=(zd, 1.0))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor((zd,))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor((zd, 1, 2))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor((zd, -1))
with self.assertRaisesRegex(TypeError, r'should be a ZstdDict object'):
ZstdDecompressor((zd, 3))
with self.assertRaises(OverflowError):
ZstdDecompressor((zd, 2**1000))
with self.assertRaises(OverflowError):
ZstdDecompressor((zd, -2**1000))
def test_train_dict(self):
TRAINED_DICT = train_dict(SAMPLES, DICT_SIZE1)
ZstdDict(TRAINED_DICT.dict_content, is_raw=False)
self.assertNotEqual(TRAINED_DICT.dict_id, 0)
self.assertGreater(len(TRAINED_DICT.dict_content), 0)
self.assertLessEqual(len(TRAINED_DICT.dict_content), DICT_SIZE1)
self.assertTrue(re.match(r'^<ZstdDict dict_id=\d+ dict_size=\d+>$', str(TRAINED_DICT)))
# compress/decompress
c = ZstdCompressor(zstd_dict=TRAINED_DICT)
for sample in SAMPLES:
dat1 = compress(sample, zstd_dict=TRAINED_DICT)
dat2 = decompress(dat1, TRAINED_DICT)
self.assertEqual(sample, dat2)
dat1 = c.compress(sample)
dat1 += c.flush()
dat2 = decompress(dat1, TRAINED_DICT)
self.assertEqual(sample, dat2)
def test_finalize_dict(self):
DICT_SIZE2 = 200*_1K
C_LEVEL = 6
try:
dic2 = finalize_dict(TRAINED_DICT, SAMPLES, DICT_SIZE2, C_LEVEL)
except NotImplementedError:
# < v1.4.5 at compile-time, >= v.1.4.5 at run-time
return
self.assertNotEqual(dic2.dict_id, 0)
self.assertGreater(len(dic2.dict_content), 0)
self.assertLessEqual(len(dic2.dict_content), DICT_SIZE2)
# compress/decompress
c = ZstdCompressor(C_LEVEL, zstd_dict=dic2)
for sample in SAMPLES:
dat1 = compress(sample, C_LEVEL, zstd_dict=dic2)
dat2 = decompress(dat1, dic2)
self.assertEqual(sample, dat2)
dat1 = c.compress(sample)
dat1 += c.flush()
dat2 = decompress(dat1, dic2)
self.assertEqual(sample, dat2)
# dict mismatch
self.assertNotEqual(TRAINED_DICT.dict_id, dic2.dict_id)
dat1 = compress(SAMPLES[0], zstd_dict=TRAINED_DICT)
with self.assertRaises(ZstdError):
decompress(dat1, dic2)
def test_train_dict_arguments(self):
with self.assertRaises(ValueError):
train_dict([], 100*_1K)
with self.assertRaises(ValueError):
train_dict(SAMPLES, -100)
with self.assertRaises(ValueError):
train_dict(SAMPLES, 0)
def test_finalize_dict_arguments(self):
with self.assertRaises(TypeError):
finalize_dict({1:2}, (b'aaa', b'bbb'), 100*_1K, 2)
with self.assertRaises(ValueError):
finalize_dict(TRAINED_DICT, [], 100*_1K, 2)
with self.assertRaises(ValueError):
finalize_dict(TRAINED_DICT, SAMPLES, -100, 2)
with self.assertRaises(ValueError):
finalize_dict(TRAINED_DICT, SAMPLES, 0, 2)
def test_train_dict_c(self):
# argument wrong type
with self.assertRaises(TypeError):
_zstd.train_dict({}, (), 100)
with self.assertRaises(TypeError):
_zstd.train_dict(bytearray(), (), 100)
with self.assertRaises(TypeError):
_zstd.train_dict(b'', 99, 100)
with self.assertRaises(TypeError):
_zstd.train_dict(b'', [], 100)
with self.assertRaises(TypeError):
_zstd.train_dict(b'', (), 100.1)
with self.assertRaises(TypeError):
_zstd.train_dict(b'', (99.1,), 100)
with self.assertRaises(ValueError):
_zstd.train_dict(b'abc', (4, -1), 100)
with self.assertRaises(ValueError):
_zstd.train_dict(b'abc', (2,), 100)
with self.assertRaises(ValueError):
_zstd.train_dict(b'', (99,), 100)
# size > size_t
with self.assertRaises(ValueError):
_zstd.train_dict(b'', (2**1000,), 100)
with self.assertRaises(ValueError):
_zstd.train_dict(b'', (-2**1000,), 100)
# dict_size <= 0
with self.assertRaises(ValueError):
_zstd.train_dict(b'', (), 0)
with self.assertRaises(ValueError):
_zstd.train_dict(b'', (), -1)
with self.assertRaises(ZstdError):
_zstd.train_dict(b'', (), 1)
def test_finalize_dict_c(self):
with self.assertRaises(TypeError):
_zstd.finalize_dict(1, 2, 3, 4, 5)
# argument wrong type
with self.assertRaises(TypeError):
_zstd.finalize_dict({}, b'', (), 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(bytearray(TRAINED_DICT.dict_content), b'', (), 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, {}, (), 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, bytearray(), (), 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', 99, 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', [], 100, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 100.1, 5)
with self.assertRaises(TypeError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 100, 5.1)
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'abc', (4, -1), 100, 5)
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'abc', (2,), 100, 5)
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (99,), 100, 5)
# size > size_t
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (2**1000,), 100, 5)
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (-2**1000,), 100, 5)
# dict_size <= 0
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 0, 5)
with self.assertRaises(ValueError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), -1, 5)
with self.assertRaises(OverflowError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 2**1000, 5)
with self.assertRaises(OverflowError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), -2**1000, 5)
with self.assertRaises(OverflowError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 100, 2**1000)
with self.assertRaises(OverflowError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 100, -2**1000)
with self.assertRaises(ZstdError):
_zstd.finalize_dict(TRAINED_DICT.dict_content, b'', (), 100, 5)
def test_train_buffer_protocol_samples(self):
def _nbytes(dat):
if isinstance(dat, (bytes, bytearray)):
return len(dat)
return memoryview(dat).nbytes
# prepare samples
chunk_lst = []
wrong_size_lst = []
correct_size_lst = []
for _ in range(300):
arr = array.array('Q', [random.randint(0, 20) for i in range(20)])
chunk_lst.append(arr)
correct_size_lst.append(_nbytes(arr))
wrong_size_lst.append(len(arr))
concatenation = b''.join(chunk_lst)
# wrong size list
with self.assertRaisesRegex(ValueError,
"The samples size tuple doesn't match the concatenation's size"):
_zstd.train_dict(concatenation, tuple(wrong_size_lst), 100*_1K)
# correct size list
_zstd.train_dict(concatenation, tuple(correct_size_lst), 3*_1K)
# wrong size list
with self.assertRaisesRegex(ValueError,
"The samples size tuple doesn't match the concatenation's size"):
_zstd.finalize_dict(TRAINED_DICT.dict_content,
concatenation, tuple(wrong_size_lst), 300*_1K, 5)
# correct size list
_zstd.finalize_dict(TRAINED_DICT.dict_content,
concatenation, tuple(correct_size_lst), 300*_1K, 5)
def test_as_prefix(self):
# V1
V1 = THIS_FILE_BYTES
zd = ZstdDict(V1, is_raw=True)
# V2
mid = len(V1) // 2
V2 = V1[:mid] + \
(b'a' if V1[mid] != int.from_bytes(b'a') else b'b') + \
V1[mid+1:]
# compress
dat = compress(V2, zstd_dict=zd.as_prefix)
self.assertEqual(get_frame_info(dat).dictionary_id, 0)
# decompress
self.assertEqual(decompress(dat, zd.as_prefix), V2)
# use wrong prefix
zd2 = ZstdDict(SAMPLES[0], is_raw=True)
try:
decompressed = decompress(dat, zd2.as_prefix)
except ZstdError: # expected
pass
else:
self.assertNotEqual(decompressed, V2)
# read only attribute
with self.assertRaises(AttributeError):
zd.as_prefix = b'1234'
def test_as_digested_dict(self):
zd = TRAINED_DICT
# test .as_digested_dict
dat = compress(SAMPLES[0], zstd_dict=zd.as_digested_dict)
self.assertEqual(decompress(dat, zd.as_digested_dict), SAMPLES[0])
with self.assertRaises(AttributeError):
zd.as_digested_dict = b'1234'
# test .as_undigested_dict
dat = compress(SAMPLES[0], zstd_dict=zd.as_undigested_dict)
self.assertEqual(decompress(dat, zd.as_undigested_dict), SAMPLES[0])
with self.assertRaises(AttributeError):
zd.as_undigested_dict = b'1234'
def test_advanced_compression_parameters(self):
options = {CompressionParameter.compression_level: 6,
CompressionParameter.window_log: 20,
CompressionParameter.enable_long_distance_matching: 1}
# automatically select
dat = compress(SAMPLES[0], options=options, zstd_dict=TRAINED_DICT)
self.assertEqual(decompress(dat, TRAINED_DICT), SAMPLES[0])
# explicitly select
dat = compress(SAMPLES[0], options=options, zstd_dict=TRAINED_DICT.as_digested_dict)
self.assertEqual(decompress(dat, TRAINED_DICT), SAMPLES[0])
def test_len(self):
self.assertEqual(len(TRAINED_DICT), len(TRAINED_DICT.dict_content))
self.assertIn(str(len(TRAINED_DICT)), str(TRAINED_DICT))
class FileTestCase(unittest.TestCase):
def setUp(self):
self.DECOMPRESSED_42 = b'a'*42
self.FRAME_42 = compress(self.DECOMPRESSED_42)
def test_init(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
pass
with ZstdFile(io.BytesIO(), "w") as f:
pass
with ZstdFile(io.BytesIO(), "x") as f:
pass
with ZstdFile(io.BytesIO(), "a") as f:
pass
with ZstdFile(io.BytesIO(), "w", level=12) as f:
pass
with ZstdFile(io.BytesIO(), "w", options={CompressionParameter.checksum_flag:1}) as f:
pass
with ZstdFile(io.BytesIO(), "w", options={}) as f:
pass
with ZstdFile(io.BytesIO(), "w", level=20, zstd_dict=TRAINED_DICT) as f:
pass
with ZstdFile(io.BytesIO(), "r", options={DecompressionParameter.window_log_max:25}) as f:
pass
with ZstdFile(io.BytesIO(), "r", options={}, zstd_dict=TRAINED_DICT) as f:
pass
def test_init_with_PathLike_filename(self):
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
filename = pathlib.Path(tmp_f.name)
with ZstdFile(filename, "a") as f:
f.write(DECOMPRESSED_100_PLUS_32KB)
with ZstdFile(filename) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
with ZstdFile(filename, "a") as f:
f.write(DECOMPRESSED_100_PLUS_32KB)
with ZstdFile(filename) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB * 2)
os.remove(filename)
def test_init_with_filename(self):
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
filename = pathlib.Path(tmp_f.name)
with ZstdFile(filename) as f:
pass
with ZstdFile(filename, "w") as f:
pass
with ZstdFile(filename, "a") as f:
pass
os.remove(filename)
def test_init_mode(self):
bi = io.BytesIO()
with ZstdFile(bi, "r"):
pass
with ZstdFile(bi, "rb"):
pass
with ZstdFile(bi, "w"):
pass
with ZstdFile(bi, "wb"):
pass
with ZstdFile(bi, "a"):
pass
with ZstdFile(bi, "ab"):
pass
def test_init_with_x_mode(self):
with tempfile.NamedTemporaryFile() as tmp_f:
filename = pathlib.Path(tmp_f.name)
for mode in ("x", "xb"):
with ZstdFile(filename, mode):
pass
with self.assertRaises(FileExistsError):
with ZstdFile(filename, mode):
pass
os.remove(filename)
def test_init_bad_mode(self):
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), (3, "x"))
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "xt")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "x+")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "rx")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "wx")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "rt")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "r+")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "wt")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "w+")
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "rw")
with self.assertRaisesRegex(TypeError,
r"not be a CompressionParameter"):
ZstdFile(io.BytesIO(), 'rb',
options={CompressionParameter.compression_level:5})
with self.assertRaisesRegex(TypeError,
r"not be a DecompressionParameter"):
ZstdFile(io.BytesIO(), 'wb',
options={DecompressionParameter.window_log_max:21})
with self.assertRaises(TypeError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "r", level=12)
def test_init_bad_check(self):
with self.assertRaises(TypeError):
ZstdFile(io.BytesIO(), "w", level='asd')
# CHECK_UNKNOWN and anything above CHECK_ID_MAX should be invalid.
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(), "w", options={999:9999})
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(), "w", options={CompressionParameter.window_log:99})
with self.assertRaises(TypeError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "r", options=33)
with self.assertRaises(OverflowError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB),
options={DecompressionParameter.window_log_max:2**31})
with self.assertRaises(ValueError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB),
options={444:333})
with self.assertRaises(TypeError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), zstd_dict={1:2})
with self.assertRaises(TypeError):
ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), zstd_dict=b'dict123456')
def test_init_close_fp(self):
# get a temp file name
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
tmp_f.write(DAT_130K_C)
filename = tmp_f.name
with self.assertRaises(TypeError):
ZstdFile(filename, options={'a':'b'})
# for PyPy
gc.collect()
os.remove(filename)
def test_close(self):
with io.BytesIO(COMPRESSED_100_PLUS_32KB) as src:
f = ZstdFile(src)
f.close()
# ZstdFile.close() should not close the underlying file object.
self.assertFalse(src.closed)
# Try closing an already-closed ZstdFile.
f.close()
self.assertFalse(src.closed)
# Test with a real file on disk, opened directly by ZstdFile.
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
filename = pathlib.Path(tmp_f.name)
f = ZstdFile(filename)
fp = f._fp
f.close()
# Here, ZstdFile.close() *should* close the underlying file object.
self.assertTrue(fp.closed)
# Try closing an already-closed ZstdFile.
f.close()
os.remove(filename)
def test_closed(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
self.assertFalse(f.closed)
f.read()
self.assertFalse(f.closed)
finally:
f.close()
self.assertTrue(f.closed)
f = ZstdFile(io.BytesIO(), "w")
try:
self.assertFalse(f.closed)
finally:
f.close()
self.assertTrue(f.closed)
def test_fileno(self):
# 1
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
self.assertRaises(io.UnsupportedOperation, f.fileno)
finally:
f.close()
self.assertRaises(ValueError, f.fileno)
# 2
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
filename = pathlib.Path(tmp_f.name)
f = ZstdFile(filename)
try:
self.assertEqual(f.fileno(), f._fp.fileno())
self.assertIsInstance(f.fileno(), int)
finally:
f.close()
self.assertRaises(ValueError, f.fileno)
os.remove(filename)
# 3, no .fileno() method
class C:
def read(self, size=-1):
return b'123'
with ZstdFile(C(), 'rb') as f:
with self.assertRaisesRegex(AttributeError, r'fileno'):
f.fileno()
def test_name(self):
# 1
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
with self.assertRaises(AttributeError):
f.name
finally:
f.close()
with self.assertRaises(ValueError):
f.name
# 2
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
filename = pathlib.Path(tmp_f.name)
f = ZstdFile(filename)
try:
self.assertEqual(f.name, f._fp.name)
self.assertIsInstance(f.name, str)
finally:
f.close()
with self.assertRaises(ValueError):
f.name
os.remove(filename)
# 3, no .filename property
class C:
def read(self, size=-1):
return b'123'
with ZstdFile(C(), 'rb') as f:
with self.assertRaisesRegex(AttributeError, r'name'):
f.name
def test_seekable(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
self.assertTrue(f.seekable())
f.read()
self.assertTrue(f.seekable())
finally:
f.close()
self.assertRaises(ValueError, f.seekable)
f = ZstdFile(io.BytesIO(), "w")
try:
self.assertFalse(f.seekable())
finally:
f.close()
self.assertRaises(ValueError, f.seekable)
src = io.BytesIO(COMPRESSED_100_PLUS_32KB)
src.seekable = lambda: False
f = ZstdFile(src)
try:
self.assertFalse(f.seekable())
finally:
f.close()
self.assertRaises(ValueError, f.seekable)
def test_readable(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
self.assertTrue(f.readable())
f.read()
self.assertTrue(f.readable())
finally:
f.close()
self.assertRaises(ValueError, f.readable)
f = ZstdFile(io.BytesIO(), "w")
try:
self.assertFalse(f.readable())
finally:
f.close()
self.assertRaises(ValueError, f.readable)
def test_writable(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
try:
self.assertFalse(f.writable())
f.read()
self.assertFalse(f.writable())
finally:
f.close()
self.assertRaises(ValueError, f.writable)
f = ZstdFile(io.BytesIO(), "w")
try:
self.assertTrue(f.writable())
finally:
f.close()
self.assertRaises(ValueError, f.writable)
def test_read_0(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
self.assertEqual(f.read(0), b"")
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB),
options={DecompressionParameter.window_log_max:20}) as f:
self.assertEqual(f.read(0), b"")
# empty file
with ZstdFile(io.BytesIO(b'')) as f:
self.assertEqual(f.read(0), b"")
with self.assertRaises(EOFError):
f.read(10)
with ZstdFile(io.BytesIO(b'')) as f:
with self.assertRaises(EOFError):
f.read(10)
def test_read_10(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
chunks = []
while True:
result = f.read(10)
if not result:
break
self.assertLessEqual(len(result), 10)
chunks.append(result)
self.assertEqual(b"".join(chunks), DECOMPRESSED_100_PLUS_32KB)
def test_read_multistream(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB * 5)) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB * 5)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB + SKIPPABLE_FRAME)) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB + COMPRESSED_DAT)) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB + DECOMPRESSED_DAT)
def test_read_incomplete(self):
with ZstdFile(io.BytesIO(DAT_130K_C[:-200])) as f:
self.assertRaises(EOFError, f.read)
# Trailing data isn't a valid compressed stream
with ZstdFile(io.BytesIO(self.FRAME_42 + b'12345')) as f:
self.assertRaises(ZstdError, f.read)
with ZstdFile(io.BytesIO(SKIPPABLE_FRAME + b'12345')) as f:
self.assertRaises(ZstdError, f.read)
def test_read_truncated(self):
# Drop stream epilogue: 4 bytes checksum
truncated = DAT_130K_C[:-4]
with ZstdFile(io.BytesIO(truncated)) as f:
self.assertRaises(EOFError, f.read)
with ZstdFile(io.BytesIO(truncated)) as f:
# this is an important test, make sure it doesn't raise EOFError.
self.assertEqual(f.read(130*_1K), DAT_130K_D)
with self.assertRaises(EOFError):
f.read(1)
# Incomplete header
for i in range(1, 20):
with ZstdFile(io.BytesIO(truncated[:i])) as f:
self.assertRaises(EOFError, f.read, 1)
def test_read_bad_args(self):
f = ZstdFile(io.BytesIO(COMPRESSED_DAT))
f.close()
self.assertRaises(ValueError, f.read)
with ZstdFile(io.BytesIO(), "w") as f:
self.assertRaises(ValueError, f.read)
with ZstdFile(io.BytesIO(COMPRESSED_DAT)) as f:
self.assertRaises(TypeError, f.read, float())
def test_read_bad_data(self):
with ZstdFile(io.BytesIO(COMPRESSED_BOGUS)) as f:
self.assertRaises(ZstdError, f.read)
def test_read_exception(self):
class C:
def read(self, size=-1):
raise OSError
with ZstdFile(C()) as f:
with self.assertRaises(OSError):
f.read(10)
def test_read1(self):
with ZstdFile(io.BytesIO(DAT_130K_C)) as f:
blocks = []
while True:
result = f.read1()
if not result:
break
blocks.append(result)
self.assertEqual(b"".join(blocks), DAT_130K_D)
self.assertEqual(f.read1(), b"")
def test_read1_0(self):
with ZstdFile(io.BytesIO(COMPRESSED_DAT)) as f:
self.assertEqual(f.read1(0), b"")
def test_read1_10(self):
with ZstdFile(io.BytesIO(COMPRESSED_DAT)) as f:
blocks = []
while True:
result = f.read1(10)
if not result:
break
blocks.append(result)
self.assertEqual(b"".join(blocks), DECOMPRESSED_DAT)
self.assertEqual(f.read1(), b"")
def test_read1_multistream(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB * 5)) as f:
blocks = []
while True:
result = f.read1()
if not result:
break
blocks.append(result)
self.assertEqual(b"".join(blocks), DECOMPRESSED_100_PLUS_32KB * 5)
self.assertEqual(f.read1(), b"")
def test_read1_bad_args(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
f.close()
self.assertRaises(ValueError, f.read1)
with ZstdFile(io.BytesIO(), "w") as f:
self.assertRaises(ValueError, f.read1)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
self.assertRaises(TypeError, f.read1, None)
def test_readinto(self):
arr = array.array("I", range(100))
self.assertEqual(len(arr), 100)
self.assertEqual(len(arr) * arr.itemsize, 400)
ba = bytearray(300)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
# 0 length output buffer
self.assertEqual(f.readinto(ba[0:0]), 0)
# use correct length for buffer protocol object
self.assertEqual(f.readinto(arr), 400)
self.assertEqual(arr.tobytes(), DECOMPRESSED_100_PLUS_32KB[:400])
# normal readinto
self.assertEqual(f.readinto(ba), 300)
self.assertEqual(ba, DECOMPRESSED_100_PLUS_32KB[400:700])
def test_peek(self):
with ZstdFile(io.BytesIO(DAT_130K_C)) as f:
result = f.peek()
self.assertGreater(len(result), 0)
self.assertTrue(DAT_130K_D.startswith(result))
self.assertEqual(f.read(), DAT_130K_D)
with ZstdFile(io.BytesIO(DAT_130K_C)) as f:
result = f.peek(10)
self.assertGreater(len(result), 0)
self.assertTrue(DAT_130K_D.startswith(result))
self.assertEqual(f.read(), DAT_130K_D)
def test_peek_bad_args(self):
with ZstdFile(io.BytesIO(), "w") as f:
self.assertRaises(ValueError, f.peek)
def test_iterator(self):
with io.BytesIO(THIS_FILE_BYTES) as f:
lines = f.readlines()
compressed = compress(THIS_FILE_BYTES)
# iter
with ZstdFile(io.BytesIO(compressed)) as f:
self.assertListEqual(list(iter(f)), lines)
# readline
with ZstdFile(io.BytesIO(compressed)) as f:
for line in lines:
self.assertEqual(f.readline(), line)
self.assertEqual(f.readline(), b'')
self.assertEqual(f.readline(), b'')
# readlines
with ZstdFile(io.BytesIO(compressed)) as f:
self.assertListEqual(f.readlines(), lines)
def test_decompress_limited(self):
_ZSTD_DStreamInSize = 128*_1K + 3
bomb = compress(b'\0' * int(2e6), level=10)
self.assertLess(len(bomb), _ZSTD_DStreamInSize)
decomp = ZstdFile(io.BytesIO(bomb))
self.assertEqual(decomp.read(1), b'\0')
# BufferedReader uses 128 KiB buffer in __init__.py
max_decomp = 128*_1K
self.assertLessEqual(decomp._buffer.raw.tell(), max_decomp,
"Excessive amount of data was decompressed")
def test_write(self):
raw_data = THIS_FILE_BYTES[: len(THIS_FILE_BYTES) // 6]
with io.BytesIO() as dst:
with ZstdFile(dst, "w") as f:
f.write(raw_data)
comp = ZstdCompressor()
expected = comp.compress(raw_data) + comp.flush()
self.assertEqual(dst.getvalue(), expected)
with io.BytesIO() as dst:
with ZstdFile(dst, "w", level=12) as f:
f.write(raw_data)
comp = ZstdCompressor(12)
expected = comp.compress(raw_data) + comp.flush()
self.assertEqual(dst.getvalue(), expected)
with io.BytesIO() as dst:
with ZstdFile(dst, "w", options={CompressionParameter.checksum_flag:1}) as f:
f.write(raw_data)
comp = ZstdCompressor(options={CompressionParameter.checksum_flag:1})
expected = comp.compress(raw_data) + comp.flush()
self.assertEqual(dst.getvalue(), expected)
with io.BytesIO() as dst:
options = {CompressionParameter.compression_level:-5,
CompressionParameter.checksum_flag:1}
with ZstdFile(dst, "w",
options=options) as f:
f.write(raw_data)
comp = ZstdCompressor(options=options)
expected = comp.compress(raw_data) + comp.flush()
self.assertEqual(dst.getvalue(), expected)
def test_write_empty_frame(self):
# .FLUSH_FRAME generates an empty content frame
c = ZstdCompressor()
self.assertNotEqual(c.flush(c.FLUSH_FRAME), b'')
self.assertNotEqual(c.flush(c.FLUSH_FRAME), b'')
# don't generate empty content frame
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
pass
self.assertEqual(bo.getvalue(), b'')
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.flush(f.FLUSH_FRAME)
self.assertEqual(bo.getvalue(), b'')
# if .write(b''), generate empty content frame
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.write(b'')
self.assertNotEqual(bo.getvalue(), b'')
# has an empty content frame
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.flush(f.FLUSH_BLOCK)
self.assertNotEqual(bo.getvalue(), b'')
def test_write_empty_block(self):
# If no internal data, .FLUSH_BLOCK return b''.
c = ZstdCompressor()
self.assertEqual(c.flush(c.FLUSH_BLOCK), b'')
self.assertNotEqual(c.compress(b'123', c.FLUSH_BLOCK),
b'')
self.assertEqual(c.flush(c.FLUSH_BLOCK), b'')
self.assertEqual(c.compress(b''), b'')
self.assertEqual(c.compress(b''), b'')
self.assertEqual(c.flush(c.FLUSH_BLOCK), b'')
# mode = .last_mode
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.write(b'123')
f.flush(f.FLUSH_BLOCK)
fp_pos = f._fp.tell()
self.assertNotEqual(fp_pos, 0)
f.flush(f.FLUSH_BLOCK)
self.assertEqual(f._fp.tell(), fp_pos)
# mode != .last_mode
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.flush(f.FLUSH_BLOCK)
self.assertEqual(f._fp.tell(), 0)
f.write(b'')
f.flush(f.FLUSH_BLOCK)
self.assertEqual(f._fp.tell(), 0)
def test_write_101(self):
with io.BytesIO() as dst:
with ZstdFile(dst, "w") as f:
for start in range(0, len(THIS_FILE_BYTES), 101):
f.write(THIS_FILE_BYTES[start:start+101])
comp = ZstdCompressor()
expected = comp.compress(THIS_FILE_BYTES) + comp.flush()
self.assertEqual(dst.getvalue(), expected)
def test_write_append(self):
def comp(data):
comp = ZstdCompressor()
return comp.compress(data) + comp.flush()
part1 = THIS_FILE_BYTES[:_1K]
part2 = THIS_FILE_BYTES[_1K:1536]
part3 = THIS_FILE_BYTES[1536:]
expected = b"".join(comp(x) for x in (part1, part2, part3))
with io.BytesIO() as dst:
with ZstdFile(dst, "w") as f:
f.write(part1)
with ZstdFile(dst, "a") as f:
f.write(part2)
with ZstdFile(dst, "a") as f:
f.write(part3)
self.assertEqual(dst.getvalue(), expected)
def test_write_bad_args(self):
f = ZstdFile(io.BytesIO(), "w")
f.close()
self.assertRaises(ValueError, f.write, b"foo")
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB), "r") as f:
self.assertRaises(ValueError, f.write, b"bar")
with ZstdFile(io.BytesIO(), "w") as f:
self.assertRaises(TypeError, f.write, None)
self.assertRaises(TypeError, f.write, "text")
self.assertRaises(TypeError, f.write, 789)
def test_writelines(self):
def comp(data):
comp = ZstdCompressor()
return comp.compress(data) + comp.flush()
with io.BytesIO(THIS_FILE_BYTES) as f:
lines = f.readlines()
with io.BytesIO() as dst:
with ZstdFile(dst, "w") as f:
f.writelines(lines)
expected = comp(THIS_FILE_BYTES)
self.assertEqual(dst.getvalue(), expected)
def test_seek_forward(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.seek(555)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[555:])
def test_seek_forward_across_streams(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB * 2)) as f:
f.seek(len(DECOMPRESSED_100_PLUS_32KB) + 123)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[123:])
def test_seek_forward_relative_to_current(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.read(100)
f.seek(1236, 1)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[1336:])
def test_seek_forward_relative_to_end(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.seek(-555, 2)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[-555:])
def test_seek_backward(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.read(1001)
f.seek(211)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[211:])
def test_seek_backward_across_streams(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB * 2)) as f:
f.read(len(DECOMPRESSED_100_PLUS_32KB) + 333)
f.seek(737)
self.assertEqual(f.read(),
DECOMPRESSED_100_PLUS_32KB[737:] + DECOMPRESSED_100_PLUS_32KB)
def test_seek_backward_relative_to_end(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.seek(-150, 2)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB[-150:])
def test_seek_past_end(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.seek(len(DECOMPRESSED_100_PLUS_32KB) + 9001)
self.assertEqual(f.tell(), len(DECOMPRESSED_100_PLUS_32KB))
self.assertEqual(f.read(), b"")
def test_seek_past_start(self):
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
f.seek(-88)
self.assertEqual(f.tell(), 0)
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
def test_seek_bad_args(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
f.close()
self.assertRaises(ValueError, f.seek, 0)
with ZstdFile(io.BytesIO(), "w") as f:
self.assertRaises(ValueError, f.seek, 0)
with ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB)) as f:
self.assertRaises(ValueError, f.seek, 0, 3)
# io.BufferedReader raises TypeError instead of ValueError
self.assertRaises((TypeError, ValueError), f.seek, 9, ())
self.assertRaises(TypeError, f.seek, None)
self.assertRaises(TypeError, f.seek, b"derp")
def test_seek_not_seekable(self):
class C(io.BytesIO):
def seekable(self):
return False
obj = C(COMPRESSED_100_PLUS_32KB)
with ZstdFile(obj, 'r') as f:
d = f.read(1)
self.assertFalse(f.seekable())
with self.assertRaisesRegex(io.UnsupportedOperation,
'File or stream is not seekable'):
f.seek(0)
d += f.read()
self.assertEqual(d, DECOMPRESSED_100_PLUS_32KB)
def test_tell(self):
with ZstdFile(io.BytesIO(DAT_130K_C)) as f:
pos = 0
while True:
self.assertEqual(f.tell(), pos)
result = f.read(random.randint(171, 189))
if not result:
break
pos += len(result)
self.assertEqual(f.tell(), len(DAT_130K_D))
with ZstdFile(io.BytesIO(), "w") as f:
for pos in range(0, len(DAT_130K_D), 143):
self.assertEqual(f.tell(), pos)
f.write(DAT_130K_D[pos:pos+143])
self.assertEqual(f.tell(), len(DAT_130K_D))
def test_tell_bad_args(self):
f = ZstdFile(io.BytesIO(COMPRESSED_100_PLUS_32KB))
f.close()
self.assertRaises(ValueError, f.tell)
def test_file_dict(self):
# default
bi = io.BytesIO()
with ZstdFile(bi, 'w', zstd_dict=TRAINED_DICT) as f:
f.write(SAMPLES[0])
bi.seek(0)
with ZstdFile(bi, zstd_dict=TRAINED_DICT) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
# .as_(un)digested_dict
bi = io.BytesIO()
with ZstdFile(bi, 'w', zstd_dict=TRAINED_DICT.as_digested_dict) as f:
f.write(SAMPLES[0])
bi.seek(0)
with ZstdFile(bi, zstd_dict=TRAINED_DICT.as_undigested_dict) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
def test_file_prefix(self):
bi = io.BytesIO()
with ZstdFile(bi, 'w', zstd_dict=TRAINED_DICT.as_prefix) as f:
f.write(SAMPLES[0])
bi.seek(0)
with ZstdFile(bi, zstd_dict=TRAINED_DICT.as_prefix) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
def test_UnsupportedOperation(self):
# 1
with ZstdFile(io.BytesIO(), 'r') as f:
with self.assertRaises(io.UnsupportedOperation):
f.write(b'1234')
# 2
class T:
def read(self, size):
return b'a' * size
with self.assertRaises(TypeError): # on creation
with ZstdFile(T(), 'w') as f:
pass
# 3
with ZstdFile(io.BytesIO(), 'w') as f:
with self.assertRaises(io.UnsupportedOperation):
f.read(100)
with self.assertRaises(io.UnsupportedOperation):
f.seek(100)
with self.assertRaises(io.UnsupportedOperation):
next(iter(f))
self.assertEqual(f.closed, True)
with self.assertRaises(ValueError):
f.readable()
with self.assertRaises(ValueError):
f.tell()
with self.assertRaises(ValueError):
f.read(100)
def test_read_readinto_readinto1(self):
lst = []
with ZstdFile(io.BytesIO(COMPRESSED_THIS_FILE*5)) as f:
while True:
method = random.randint(0, 2)
size = random.randint(0, 300)
if method == 0:
dat = f.read(size)
if not dat and size:
break
lst.append(dat)
elif method == 1:
ba = bytearray(size)
read_size = f.readinto(ba)
if read_size == 0 and size:
break
lst.append(bytes(ba[:read_size]))
elif method == 2:
ba = bytearray(size)
read_size = f.readinto1(ba)
if read_size == 0 and size:
break
lst.append(bytes(ba[:read_size]))
self.assertEqual(b''.join(lst), THIS_FILE_BYTES*5)
def test_zstdfile_flush(self):
# closed
f = ZstdFile(io.BytesIO(), 'w')
f.close()
with self.assertRaises(ValueError):
f.flush()
# read
with ZstdFile(io.BytesIO(), 'r') as f:
# does nothing for read-only stream
f.flush()
# write
DAT = b'abcd'
bi = io.BytesIO()
with ZstdFile(bi, 'w') as f:
self.assertEqual(f.write(DAT), len(DAT))
self.assertEqual(f.tell(), len(DAT))
self.assertEqual(bi.tell(), 0) # not enough for a block
self.assertEqual(f.flush(), None)
self.assertEqual(f.tell(), len(DAT))
self.assertGreater(bi.tell(), 0) # flushed
# write, no .flush() method
class C:
def write(self, b):
return len(b)
with ZstdFile(C(), 'w') as f:
self.assertEqual(f.write(DAT), len(DAT))
self.assertEqual(f.tell(), len(DAT))
self.assertEqual(f.flush(), None)
self.assertEqual(f.tell(), len(DAT))
def test_zstdfile_flush_mode(self):
self.assertEqual(ZstdFile.FLUSH_BLOCK, ZstdCompressor.FLUSH_BLOCK)
self.assertEqual(ZstdFile.FLUSH_FRAME, ZstdCompressor.FLUSH_FRAME)
with self.assertRaises(AttributeError):
ZstdFile.CONTINUE
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
# flush block
self.assertEqual(f.write(b'123'), 3)
self.assertIsNone(f.flush(f.FLUSH_BLOCK))
p1 = bo.tell()
# mode == .last_mode, should return
self.assertIsNone(f.flush())
p2 = bo.tell()
self.assertEqual(p1, p2)
# flush frame
self.assertEqual(f.write(b'456'), 3)
self.assertIsNone(f.flush(mode=f.FLUSH_FRAME))
# flush frame
self.assertEqual(f.write(b'789'), 3)
self.assertIsNone(f.flush(f.FLUSH_FRAME))
p1 = bo.tell()
# mode == .last_mode, should return
self.assertIsNone(f.flush(f.FLUSH_FRAME))
p2 = bo.tell()
self.assertEqual(p1, p2)
self.assertEqual(decompress(bo.getvalue()), b'123456789')
bo = io.BytesIO()
with ZstdFile(bo, 'w') as f:
f.write(b'123')
with self.assertRaisesRegex(ValueError, r'\.FLUSH_.*?\.FLUSH_'):
f.flush(ZstdCompressor.CONTINUE)
with self.assertRaises(ValueError):
f.flush(-1)
with self.assertRaises(ValueError):
f.flush(123456)
with self.assertRaises(TypeError):
f.flush(node=ZstdCompressor.CONTINUE)
with self.assertRaises((TypeError, ValueError)):
f.flush('FLUSH_FRAME')
with self.assertRaises(TypeError):
f.flush(b'456', f.FLUSH_BLOCK)
def test_zstdfile_truncate(self):
with ZstdFile(io.BytesIO(), 'w') as f:
with self.assertRaises(io.UnsupportedOperation):
f.truncate(200)
def test_zstdfile_iter_issue45475(self):
lines = [l for l in ZstdFile(io.BytesIO(COMPRESSED_THIS_FILE))]
self.assertGreater(len(lines), 0)
def test_append_new_file(self):
with tempfile.NamedTemporaryFile(delete=True) as tmp_f:
filename = tmp_f.name
with ZstdFile(filename, 'a') as f:
pass
self.assertTrue(os.path.isfile(filename))
os.remove(filename)
class OpenTestCase(unittest.TestCase):
def test_binary_modes(self):
with open(io.BytesIO(COMPRESSED_100_PLUS_32KB), "rb") as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
with io.BytesIO() as bio:
with open(bio, "wb") as f:
f.write(DECOMPRESSED_100_PLUS_32KB)
file_data = decompress(bio.getvalue())
self.assertEqual(file_data, DECOMPRESSED_100_PLUS_32KB)
with open(bio, "ab") as f:
f.write(DECOMPRESSED_100_PLUS_32KB)
file_data = decompress(bio.getvalue())
self.assertEqual(file_data, DECOMPRESSED_100_PLUS_32KB * 2)
def test_text_modes(self):
# empty input
with self.assertRaises(EOFError):
with open(io.BytesIO(b''), "rt", encoding="utf-8", newline='\n') as reader:
for _ in reader:
pass
# read
uncompressed = THIS_FILE_STR.replace(os.linesep, "\n")
with open(io.BytesIO(COMPRESSED_THIS_FILE), "rt", encoding="utf-8") as f:
self.assertEqual(f.read(), uncompressed)
with io.BytesIO() as bio:
# write
with open(bio, "wt", encoding="utf-8") as f:
f.write(uncompressed)
file_data = decompress(bio.getvalue()).decode("utf-8")
self.assertEqual(file_data.replace(os.linesep, "\n"), uncompressed)
# append
with open(bio, "at", encoding="utf-8") as f:
f.write(uncompressed)
file_data = decompress(bio.getvalue()).decode("utf-8")
self.assertEqual(file_data.replace(os.linesep, "\n"), uncompressed * 2)
def test_bad_params(self):
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
TESTFN = pathlib.Path(tmp_f.name)
with self.assertRaises(ValueError):
open(TESTFN, "")
with self.assertRaises(ValueError):
open(TESTFN, "rbt")
with self.assertRaises(ValueError):
open(TESTFN, "rb", encoding="utf-8")
with self.assertRaises(ValueError):
open(TESTFN, "rb", errors="ignore")
with self.assertRaises(ValueError):
open(TESTFN, "rb", newline="\n")
os.remove(TESTFN)
def test_option(self):
options = {DecompressionParameter.window_log_max:25}
with open(io.BytesIO(COMPRESSED_100_PLUS_32KB), "rb", options=options) as f:
self.assertEqual(f.read(), DECOMPRESSED_100_PLUS_32KB)
options = {CompressionParameter.compression_level:12}
with io.BytesIO() as bio:
with open(bio, "wb", options=options) as f:
f.write(DECOMPRESSED_100_PLUS_32KB)
file_data = decompress(bio.getvalue())
self.assertEqual(file_data, DECOMPRESSED_100_PLUS_32KB)
def test_encoding(self):
uncompressed = THIS_FILE_STR.replace(os.linesep, "\n")
with io.BytesIO() as bio:
with open(bio, "wt", encoding="utf-16-le") as f:
f.write(uncompressed)
file_data = decompress(bio.getvalue()).decode("utf-16-le")
self.assertEqual(file_data.replace(os.linesep, "\n"), uncompressed)
bio.seek(0)
with open(bio, "rt", encoding="utf-16-le") as f:
self.assertEqual(f.read().replace(os.linesep, "\n"), uncompressed)
def test_encoding_error_handler(self):
with io.BytesIO(compress(b"foo\xffbar")) as bio:
with open(bio, "rt", encoding="ascii", errors="ignore") as f:
self.assertEqual(f.read(), "foobar")
def test_newline(self):
# Test with explicit newline (universal newline mode disabled).
text = THIS_FILE_STR.replace(os.linesep, "\n")
with io.BytesIO() as bio:
with open(bio, "wt", encoding="utf-8", newline="\n") as f:
f.write(text)
bio.seek(0)
with open(bio, "rt", encoding="utf-8", newline="\r") as f:
self.assertEqual(f.readlines(), [text])
def test_x_mode(self):
with tempfile.NamedTemporaryFile(delete=False) as tmp_f:
TESTFN = pathlib.Path(tmp_f.name)
for mode in ("x", "xb", "xt"):
os.remove(TESTFN)
if mode == "xt":
encoding = "utf-8"
else:
encoding = None
with open(TESTFN, mode, encoding=encoding):
pass
with self.assertRaises(FileExistsError):
with open(TESTFN, mode):
pass
os.remove(TESTFN)
def test_open_dict(self):
# default
bi = io.BytesIO()
with open(bi, 'w', zstd_dict=TRAINED_DICT) as f:
f.write(SAMPLES[0])
bi.seek(0)
with open(bi, zstd_dict=TRAINED_DICT) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
# .as_(un)digested_dict
bi = io.BytesIO()
with open(bi, 'w', zstd_dict=TRAINED_DICT.as_digested_dict) as f:
f.write(SAMPLES[0])
bi.seek(0)
with open(bi, zstd_dict=TRAINED_DICT.as_undigested_dict) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
# invalid dictionary
bi = io.BytesIO()
with self.assertRaisesRegex(TypeError, 'zstd_dict'):
open(bi, 'w', zstd_dict={1:2, 2:3})
with self.assertRaisesRegex(TypeError, 'zstd_dict'):
open(bi, 'w', zstd_dict=b'1234567890')
def test_open_prefix(self):
bi = io.BytesIO()
with open(bi, 'w', zstd_dict=TRAINED_DICT.as_prefix) as f:
f.write(SAMPLES[0])
bi.seek(0)
with open(bi, zstd_dict=TRAINED_DICT.as_prefix) as f:
dat = f.read()
self.assertEqual(dat, SAMPLES[0])
def test_buffer_protocol(self):
# don't use len() for buffer protocol objects
arr = array.array("i", range(1000))
LENGTH = len(arr) * arr.itemsize
with open(io.BytesIO(), "wb") as f:
self.assertEqual(f.write(arr), LENGTH)
self.assertEqual(f.tell(), LENGTH)
class FreeThreadingMethodTests(unittest.TestCase):
@threading_helper.reap_threads
@threading_helper.requires_working_threading()
def test_compress_locking(self):
input = b'a'* (16*_1K)
num_threads = 8
# gh-136394: the first output of .compress() includes the frame header
# we run the first .compress() call outside of the threaded portion
# to make the test order-independent
comp = ZstdCompressor()
parts = [comp.compress(input, ZstdCompressor.FLUSH_BLOCK)]
for _ in range(num_threads):
res = comp.compress(input, ZstdCompressor.FLUSH_BLOCK)
if res:
parts.append(res)
rest1 = comp.flush()
expected = b''.join(parts) + rest1
comp = ZstdCompressor()
output = [comp.compress(input, ZstdCompressor.FLUSH_BLOCK)]
def run_method(method, input_data, output_data):
res = method(input_data, ZstdCompressor.FLUSH_BLOCK)
if res:
output_data.append(res)
threads = []
for i in range(num_threads):
thread = threading.Thread(target=run_method, args=(comp.compress, input, output))
threads.append(thread)
with threading_helper.start_threads(threads):
pass
rest2 = comp.flush()
self.assertEqual(rest1, rest2)
actual = b''.join(output) + rest2
self.assertEqual(expected, actual)
@threading_helper.reap_threads
@threading_helper.requires_working_threading()
def test_decompress_locking(self):
input = compress(b'a'* (16*_1K))
num_threads = 8
# to ensure we decompress over multiple calls, set maxsize
window_size = _1K * 16//num_threads
decomp = ZstdDecompressor()
parts = []
for _ in range(num_threads):
res = decomp.decompress(input, window_size)
if res:
parts.append(res)
expected = b''.join(parts)
comp = ZstdDecompressor()
output = []
def run_method(method, input_data, output_data):
res = method(input_data, window_size)
if res:
output_data.append(res)
threads = []
for i in range(num_threads):
thread = threading.Thread(target=run_method, args=(comp.decompress, input, output))
threads.append(thread)
with threading_helper.start_threads(threads):
pass
actual = b''.join(output)
self.assertEqual(expected, actual)
@threading_helper.reap_threads
@threading_helper.requires_working_threading()
def test_compress_shared_dict(self):
num_threads = 8
def run_method(b):
level = threading.get_ident() % 4
# sync threads to increase chance of contention on
# capsule storing dictionary levels
b.wait()
ZstdCompressor(level=level,
zstd_dict=TRAINED_DICT.as_digested_dict)
b.wait()
ZstdCompressor(level=level,
zstd_dict=TRAINED_DICT.as_undigested_dict)
b.wait()
ZstdCompressor(level=level,
zstd_dict=TRAINED_DICT.as_prefix)
threads = []
b = threading.Barrier(num_threads)
for i in range(num_threads):
thread = threading.Thread(target=run_method, args=(b,))
threads.append(thread)
with threading_helper.start_threads(threads):
pass
@threading_helper.reap_threads
@threading_helper.requires_working_threading()
def test_decompress_shared_dict(self):
num_threads = 8
def run_method(b):
# sync threads to increase chance of contention on
# decompression dictionary
b.wait()
ZstdDecompressor(zstd_dict=TRAINED_DICT.as_digested_dict)
b.wait()
ZstdDecompressor(zstd_dict=TRAINED_DICT.as_undigested_dict)
b.wait()
ZstdDecompressor(zstd_dict=TRAINED_DICT.as_prefix)
threads = []
b = threading.Barrier(num_threads)
for i in range(num_threads):
thread = threading.Thread(target=run_method, args=(b,))
threads.append(thread)
with threading_helper.start_threads(threads):
pass
if __name__ == "__main__":
unittest.main()