##// END OF EJS Templates
python-zstandard: blacken at 80 characters...
Gregory Szorc -
r44605:5e84a96d default
parent child Browse files
Show More
@@ -9,7 +9,6 b' build/'
9 | \.mypy_cache/
9 | \.mypy_cache/
10 | \.venv/
10 | \.venv/
11 | mercurial/thirdparty/
11 | mercurial/thirdparty/
12 | contrib/python-zstandard/
13 '''
12 '''
14 skip-string-normalization = true
13 skip-string-normalization = true
15 quiet = true
14 quiet = true
@@ -6,7 +6,7 b' rustfmt:command = rustfmt +nightly'
6 rustfmt:pattern = set:**.rs
6 rustfmt:pattern = set:**.rs
7
7
8 black:command = black --config=black.toml -
8 black:command = black --config=black.toml -
9 black:pattern = set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"
9 black:pattern = set:**.py - mercurial/thirdparty/**
10
10
11 # Mercurial doesn't have any Go code, but if we did this is how we
11 # Mercurial doesn't have any Go code, but if we did this is how we
12 # would configure `hg fix` for Go:
12 # would configure `hg fix` for Go:
@@ -52,7 +52,8 b' SOURCES = ['
52
52
53 # Headers whose preprocessed output will be fed into cdef().
53 # Headers whose preprocessed output will be fed into cdef().
54 HEADERS = [
54 HEADERS = [
55 os.path.join(HERE, "zstd", *p) for p in (("zstd.h",), ("dictBuilder", "zdict.h"),)
55 os.path.join(HERE, "zstd", *p)
56 for p in (("zstd.h",), ("dictBuilder", "zdict.h"),)
56 ]
57 ]
57
58
58 INCLUDE_DIRS = [
59 INCLUDE_DIRS = [
@@ -139,7 +140,9 b' def preprocess(path):'
139 env = dict(os.environ)
140 env = dict(os.environ)
140 if getattr(compiler, "_paths", None):
141 if getattr(compiler, "_paths", None):
141 env["PATH"] = compiler._paths
142 env["PATH"] = compiler._paths
142 process = subprocess.Popen(args + [input_file], stdout=subprocess.PIPE, env=env)
143 process = subprocess.Popen(
144 args + [input_file], stdout=subprocess.PIPE, env=env
145 )
143 output = process.communicate()[0]
146 output = process.communicate()[0]
144 ret = process.poll()
147 ret = process.poll()
145 if ret:
148 if ret:
@@ -87,7 +87,9 b' with open("c-ext/python-zstandard.h", "r'
87 break
87 break
88
88
89 if not version:
89 if not version:
90 raise Exception("could not resolve package version; " "this should never happen")
90 raise Exception(
91 "could not resolve package version; " "this should never happen"
92 )
91
93
92 setup(
94 setup(
93 name="zstandard",
95 name="zstandard",
@@ -138,12 +138,16 b' def get_c_extension('
138 if not system_zstd:
138 if not system_zstd:
139 sources.update([os.path.join(actual_root, p) for p in zstd_sources])
139 sources.update([os.path.join(actual_root, p) for p in zstd_sources])
140 if support_legacy:
140 if support_legacy:
141 sources.update([os.path.join(actual_root, p) for p in zstd_sources_legacy])
141 sources.update(
142 [os.path.join(actual_root, p) for p in zstd_sources_legacy]
143 )
142 sources = list(sources)
144 sources = list(sources)
143
145
144 include_dirs = set([os.path.join(actual_root, d) for d in ext_includes])
146 include_dirs = set([os.path.join(actual_root, d) for d in ext_includes])
145 if not system_zstd:
147 if not system_zstd:
146 include_dirs.update([os.path.join(actual_root, d) for d in zstd_includes])
148 include_dirs.update(
149 [os.path.join(actual_root, d) for d in zstd_includes]
150 )
147 if support_legacy:
151 if support_legacy:
148 include_dirs.update(
152 include_dirs.update(
149 [os.path.join(actual_root, d) for d in zstd_includes_legacy]
153 [os.path.join(actual_root, d) for d in zstd_includes_legacy]
@@ -50,7 +50,9 b' def make_cffi(cls):'
50 os.environ.update(old_env)
50 os.environ.update(old_env)
51
51
52 if mod.backend != "cffi":
52 if mod.backend != "cffi":
53 raise Exception("got the zstandard %s backend instead of cffi" % mod.backend)
53 raise Exception(
54 "got the zstandard %s backend instead of cffi" % mod.backend
55 )
54
56
55 # If CFFI version is available, dynamically construct test methods
57 # If CFFI version is available, dynamically construct test methods
56 # that use it.
58 # that use it.
@@ -84,7 +86,9 b' def make_cffi(cls):'
84 fn.__func__.func_defaults,
86 fn.__func__.func_defaults,
85 fn.__func__.func_closure,
87 fn.__func__.func_closure,
86 )
88 )
87 new_method = types.UnboundMethodType(new_fn, fn.im_self, fn.im_class)
89 new_method = types.UnboundMethodType(
90 new_fn, fn.im_self, fn.im_class
91 )
88
92
89 setattr(cls, name, new_method)
93 setattr(cls, name, new_method)
90
94
@@ -194,4 +198,6 b' if hypothesis:'
194 expensive_settings = hypothesis.settings(deadline=None, max_examples=10000)
198 expensive_settings = hypothesis.settings(deadline=None, max_examples=10000)
195 hypothesis.settings.register_profile("expensive", expensive_settings)
199 hypothesis.settings.register_profile("expensive", expensive_settings)
196
200
197 hypothesis.settings.load_profile(os.environ.get("HYPOTHESIS_PROFILE", "default"))
201 hypothesis.settings.load_profile(
202 os.environ.get("HYPOTHESIS_PROFILE", "default")
203 )
@@ -67,7 +67,8 b' class TestBufferWithSegments(TestCase):'
67 self.skipTest("BufferWithSegments not available")
67 self.skipTest("BufferWithSegments not available")
68
68
69 b = zstd.BufferWithSegments(
69 b = zstd.BufferWithSegments(
70 b"foofooxfooxy", b"".join([ss.pack(0, 3), ss.pack(3, 4), ss.pack(7, 5)])
70 b"foofooxfooxy",
71 b"".join([ss.pack(0, 3), ss.pack(3, 4), ss.pack(7, 5)]),
71 )
72 )
72 self.assertEqual(len(b), 3)
73 self.assertEqual(len(b), 3)
73 self.assertEqual(b.size, 12)
74 self.assertEqual(b.size, 12)
@@ -83,17 +84,23 b' class TestBufferWithSegmentsCollection(T'
83 if not hasattr(zstd, "BufferWithSegmentsCollection"):
84 if not hasattr(zstd, "BufferWithSegmentsCollection"):
84 self.skipTest("BufferWithSegmentsCollection not available")
85 self.skipTest("BufferWithSegmentsCollection not available")
85
86
86 with self.assertRaisesRegex(ValueError, "must pass at least 1 argument"):
87 with self.assertRaisesRegex(
88 ValueError, "must pass at least 1 argument"
89 ):
87 zstd.BufferWithSegmentsCollection()
90 zstd.BufferWithSegmentsCollection()
88
91
89 def test_argument_validation(self):
92 def test_argument_validation(self):
90 if not hasattr(zstd, "BufferWithSegmentsCollection"):
93 if not hasattr(zstd, "BufferWithSegmentsCollection"):
91 self.skipTest("BufferWithSegmentsCollection not available")
94 self.skipTest("BufferWithSegmentsCollection not available")
92
95
93 with self.assertRaisesRegex(TypeError, "arguments must be BufferWithSegments"):
96 with self.assertRaisesRegex(
97 TypeError, "arguments must be BufferWithSegments"
98 ):
94 zstd.BufferWithSegmentsCollection(None)
99 zstd.BufferWithSegmentsCollection(None)
95
100
96 with self.assertRaisesRegex(TypeError, "arguments must be BufferWithSegments"):
101 with self.assertRaisesRegex(
102 TypeError, "arguments must be BufferWithSegments"
103 ):
97 zstd.BufferWithSegmentsCollection(
104 zstd.BufferWithSegmentsCollection(
98 zstd.BufferWithSegments(b"foo", ss.pack(0, 3)), None
105 zstd.BufferWithSegments(b"foo", ss.pack(0, 3)), None
99 )
106 )
@@ -24,7 +24,9 b' else:'
24
24
25
25
26 def multithreaded_chunk_size(level, source_size=0):
26 def multithreaded_chunk_size(level, source_size=0):
27 params = zstd.ZstdCompressionParameters.from_level(level, source_size=source_size)
27 params = zstd.ZstdCompressionParameters.from_level(
28 level, source_size=source_size
29 )
28
30
29 return 1 << (params.window_log + 2)
31 return 1 << (params.window_log + 2)
30
32
@@ -86,7 +88,9 b' class TestCompressor_compress(TestCase):'
86
88
87 # This matches the test for read_to_iter() below.
89 # This matches the test for read_to_iter() below.
88 cctx = zstd.ZstdCompressor(level=1, write_content_size=False)
90 cctx = zstd.ZstdCompressor(level=1, write_content_size=False)
89 result = cctx.compress(b"f" * zstd.COMPRESSION_RECOMMENDED_INPUT_SIZE + b"o")
91 result = cctx.compress(
92 b"f" * zstd.COMPRESSION_RECOMMENDED_INPUT_SIZE + b"o"
93 )
90 self.assertEqual(
94 self.assertEqual(
91 result,
95 result,
92 b"\x28\xb5\x2f\xfd\x00\x40\x54\x00\x00"
96 b"\x28\xb5\x2f\xfd\x00\x40\x54\x00\x00"
@@ -99,7 +103,9 b' class TestCompressor_compress(TestCase):'
99 result = cctx.compress(b"foo" * 256)
103 result = cctx.compress(b"foo" * 256)
100
104
101 def test_no_magic(self):
105 def test_no_magic(self):
102 params = zstd.ZstdCompressionParameters.from_level(1, format=zstd.FORMAT_ZSTD1)
106 params = zstd.ZstdCompressionParameters.from_level(
107 1, format=zstd.FORMAT_ZSTD1
108 )
103 cctx = zstd.ZstdCompressor(compression_params=params)
109 cctx = zstd.ZstdCompressor(compression_params=params)
104 magic = cctx.compress(b"foobar")
110 magic = cctx.compress(b"foobar")
105
111
@@ -223,7 +229,8 b' class TestCompressor_compress(TestCase):'
223
229
224 self.assertEqual(
230 self.assertEqual(
225 result,
231 result,
226 b"\x28\xb5\x2f\xfd\x23\x8f\x55\x0f\x70\x03\x19\x00\x00" b"\x66\x6f\x6f",
232 b"\x28\xb5\x2f\xfd\x23\x8f\x55\x0f\x70\x03\x19\x00\x00"
233 b"\x66\x6f\x6f",
227 )
234 )
228
235
229 def test_multithreaded_compression_params(self):
236 def test_multithreaded_compression_params(self):
@@ -234,7 +241,9 b' class TestCompressor_compress(TestCase):'
234 params = zstd.get_frame_parameters(result)
241 params = zstd.get_frame_parameters(result)
235 self.assertEqual(params.content_size, 3)
242 self.assertEqual(params.content_size, 3)
236
243
237 self.assertEqual(result, b"\x28\xb5\x2f\xfd\x20\x03\x19\x00\x00\x66\x6f\x6f")
244 self.assertEqual(
245 result, b"\x28\xb5\x2f\xfd\x20\x03\x19\x00\x00\x66\x6f\x6f"
246 )
238
247
239
248
240 @make_cffi
249 @make_cffi
@@ -347,7 +356,9 b' class TestCompressor_compressobj(TestCas'
347 )
356 )
348 self.assertEqual(cobj.compress(b"bar"), b"")
357 self.assertEqual(cobj.compress(b"bar"), b"")
349 # 3 byte header plus content.
358 # 3 byte header plus content.
350 self.assertEqual(cobj.flush(zstd.COMPRESSOBJ_FLUSH_BLOCK), b"\x18\x00\x00bar")
359 self.assertEqual(
360 cobj.flush(zstd.COMPRESSOBJ_FLUSH_BLOCK), b"\x18\x00\x00bar"
361 )
351 self.assertEqual(cobj.flush(), b"\x01\x00\x00")
362 self.assertEqual(cobj.flush(), b"\x01\x00\x00")
352
363
353 def test_flush_empty_block(self):
364 def test_flush_empty_block(self):
@@ -445,7 +456,9 b' class TestCompressor_copy_stream(TestCas'
445 self.assertEqual(int(r), 0)
456 self.assertEqual(int(r), 0)
446 self.assertEqual(w, 9)
457 self.assertEqual(w, 9)
447
458
448 self.assertEqual(dest.getvalue(), b"\x28\xb5\x2f\xfd\x00\x48\x01\x00\x00")
459 self.assertEqual(
460 dest.getvalue(), b"\x28\xb5\x2f\xfd\x00\x48\x01\x00\x00"
461 )
449
462
450 def test_large_data(self):
463 def test_large_data(self):
451 source = io.BytesIO()
464 source = io.BytesIO()
@@ -478,7 +491,9 b' class TestCompressor_copy_stream(TestCas'
478 cctx = zstd.ZstdCompressor(level=1, write_checksum=True)
491 cctx = zstd.ZstdCompressor(level=1, write_checksum=True)
479 cctx.copy_stream(source, with_checksum)
492 cctx.copy_stream(source, with_checksum)
480
493
481 self.assertEqual(len(with_checksum.getvalue()), len(no_checksum.getvalue()) + 4)
494 self.assertEqual(
495 len(with_checksum.getvalue()), len(no_checksum.getvalue()) + 4
496 )
482
497
483 no_params = zstd.get_frame_parameters(no_checksum.getvalue())
498 no_params = zstd.get_frame_parameters(no_checksum.getvalue())
484 with_params = zstd.get_frame_parameters(with_checksum.getvalue())
499 with_params = zstd.get_frame_parameters(with_checksum.getvalue())
@@ -585,7 +600,9 b' class TestCompressor_stream_reader(TestC'
585 cctx = zstd.ZstdCompressor()
600 cctx = zstd.ZstdCompressor()
586
601
587 with cctx.stream_reader(b"foo") as reader:
602 with cctx.stream_reader(b"foo") as reader:
588 with self.assertRaisesRegex(ValueError, "cannot __enter__ multiple times"):
603 with self.assertRaisesRegex(
604 ValueError, "cannot __enter__ multiple times"
605 ):
589 with reader as reader2:
606 with reader as reader2:
590 pass
607 pass
591
608
@@ -744,7 +761,9 b' class TestCompressor_stream_reader(TestC'
744 source = io.BytesIO(b"foobar")
761 source = io.BytesIO(b"foobar")
745
762
746 with cctx.stream_reader(source, size=2) as reader:
763 with cctx.stream_reader(source, size=2) as reader:
747 with self.assertRaisesRegex(zstd.ZstdError, "Src size is incorrect"):
764 with self.assertRaisesRegex(
765 zstd.ZstdError, "Src size is incorrect"
766 ):
748 reader.read(10)
767 reader.read(10)
749
768
750 # Try another compression operation.
769 # Try another compression operation.
@@ -1126,7 +1145,9 b' class TestCompressor_stream_writer(TestC'
1126 self.assertFalse(no_params.has_checksum)
1145 self.assertFalse(no_params.has_checksum)
1127 self.assertTrue(with_params.has_checksum)
1146 self.assertTrue(with_params.has_checksum)
1128
1147
1129 self.assertEqual(len(with_checksum.getvalue()), len(no_checksum.getvalue()) + 4)
1148 self.assertEqual(
1149 len(with_checksum.getvalue()), len(no_checksum.getvalue()) + 4
1150 )
1130
1151
1131 def test_write_content_size(self):
1152 def test_write_content_size(self):
1132 no_size = NonClosingBytesIO()
1153 no_size = NonClosingBytesIO()
@@ -1145,7 +1166,9 b' class TestCompressor_stream_writer(TestC'
1145
1166
1146 # Declaring size will write the header.
1167 # Declaring size will write the header.
1147 with_size = NonClosingBytesIO()
1168 with_size = NonClosingBytesIO()
1148 with cctx.stream_writer(with_size, size=len(b"foobar" * 256)) as compressor:
1169 with cctx.stream_writer(
1170 with_size, size=len(b"foobar" * 256)
1171 ) as compressor:
1149 self.assertEqual(compressor.write(b"foobar" * 256), 0)
1172 self.assertEqual(compressor.write(b"foobar" * 256), 0)
1150
1173
1151 no_params = zstd.get_frame_parameters(no_size.getvalue())
1174 no_params = zstd.get_frame_parameters(no_size.getvalue())
@@ -1191,7 +1214,9 b' class TestCompressor_stream_writer(TestC'
1191 self.assertFalse(no_params.has_checksum)
1214 self.assertFalse(no_params.has_checksum)
1192 self.assertFalse(with_params.has_checksum)
1215 self.assertFalse(with_params.has_checksum)
1193
1216
1194 self.assertEqual(len(with_dict_id.getvalue()), len(no_dict_id.getvalue()) + 4)
1217 self.assertEqual(
1218 len(with_dict_id.getvalue()), len(no_dict_id.getvalue()) + 4
1219 )
1195
1220
1196 def test_memory_size(self):
1221 def test_memory_size(self):
1197 cctx = zstd.ZstdCompressor(level=3)
1222 cctx = zstd.ZstdCompressor(level=3)
@@ -1337,7 +1362,9 b' class TestCompressor_read_to_iter(TestCa'
1337 for chunk in cctx.read_to_iter(b"foobar"):
1362 for chunk in cctx.read_to_iter(b"foobar"):
1338 pass
1363 pass
1339
1364
1340 with self.assertRaisesRegex(ValueError, "must pass an object with a read"):
1365 with self.assertRaisesRegex(
1366 ValueError, "must pass an object with a read"
1367 ):
1341 for chunk in cctx.read_to_iter(True):
1368 for chunk in cctx.read_to_iter(True):
1342 pass
1369 pass
1343
1370
@@ -1513,7 +1540,9 b' class TestCompressor_chunker(TestCase):'
1513
1540
1514 dctx = zstd.ZstdDecompressor()
1541 dctx = zstd.ZstdDecompressor()
1515
1542
1516 self.assertEqual(dctx.decompress(b"".join(chunks)), (b"x" * 1000) + (b"y" * 24))
1543 self.assertEqual(
1544 dctx.decompress(b"".join(chunks)), (b"x" * 1000) + (b"y" * 24)
1545 )
1517
1546
1518 def test_small_chunk_size(self):
1547 def test_small_chunk_size(self):
1519 cctx = zstd.ZstdCompressor()
1548 cctx = zstd.ZstdCompressor()
@@ -1533,7 +1562,8 b' class TestCompressor_chunker(TestCase):'
1533
1562
1534 dctx = zstd.ZstdDecompressor()
1563 dctx = zstd.ZstdDecompressor()
1535 self.assertEqual(
1564 self.assertEqual(
1536 dctx.decompress(b"".join(chunks), max_output_size=10000), b"foo" * 1024
1565 dctx.decompress(b"".join(chunks), max_output_size=10000),
1566 b"foo" * 1024,
1537 )
1567 )
1538
1568
1539 def test_input_types(self):
1569 def test_input_types(self):
@@ -1602,7 +1632,8 b' class TestCompressor_chunker(TestCase):'
1602 list(chunker.finish())
1632 list(chunker.finish())
1603
1633
1604 with self.assertRaisesRegex(
1634 with self.assertRaisesRegex(
1605 zstd.ZstdError, r"cannot call compress\(\) after compression finished"
1635 zstd.ZstdError,
1636 r"cannot call compress\(\) after compression finished",
1606 ):
1637 ):
1607 list(chunker.compress(b"foo"))
1638 list(chunker.compress(b"foo"))
1608
1639
@@ -1644,7 +1675,9 b' class TestCompressor_multi_compress_to_b'
1644 with self.assertRaises(TypeError):
1675 with self.assertRaises(TypeError):
1645 cctx.multi_compress_to_buffer((1, 2))
1676 cctx.multi_compress_to_buffer((1, 2))
1646
1677
1647 with self.assertRaisesRegex(TypeError, "item 0 not a bytes like object"):
1678 with self.assertRaisesRegex(
1679 TypeError, "item 0 not a bytes like object"
1680 ):
1648 cctx.multi_compress_to_buffer([u"foo"])
1681 cctx.multi_compress_to_buffer([u"foo"])
1649
1682
1650 def test_empty_input(self):
1683 def test_empty_input(self):
@@ -28,9 +28,13 b' class TestCompressor_stream_reader_fuzzi'
28 original=strategies.sampled_from(random_input_data()),
28 original=strategies.sampled_from(random_input_data()),
29 level=strategies.integers(min_value=1, max_value=5),
29 level=strategies.integers(min_value=1, max_value=5),
30 source_read_size=strategies.integers(1, 16384),
30 source_read_size=strategies.integers(1, 16384),
31 read_size=strategies.integers(-1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
31 read_size=strategies.integers(
32 -1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
33 ),
32 )
34 )
33 def test_stream_source_read(self, original, level, source_read_size, read_size):
35 def test_stream_source_read(
36 self, original, level, source_read_size, read_size
37 ):
34 if read_size == 0:
38 if read_size == 0:
35 read_size = -1
39 read_size = -1
36
40
@@ -58,9 +62,13 b' class TestCompressor_stream_reader_fuzzi'
58 original=strategies.sampled_from(random_input_data()),
62 original=strategies.sampled_from(random_input_data()),
59 level=strategies.integers(min_value=1, max_value=5),
63 level=strategies.integers(min_value=1, max_value=5),
60 source_read_size=strategies.integers(1, 16384),
64 source_read_size=strategies.integers(1, 16384),
61 read_size=strategies.integers(-1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
65 read_size=strategies.integers(
66 -1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
67 ),
62 )
68 )
63 def test_buffer_source_read(self, original, level, source_read_size, read_size):
69 def test_buffer_source_read(
70 self, original, level, source_read_size, read_size
71 ):
64 if read_size == 0:
72 if read_size == 0:
65 read_size = -1
73 read_size = -1
66
74
@@ -155,9 +163,13 b' class TestCompressor_stream_reader_fuzzi'
155 original=strategies.sampled_from(random_input_data()),
163 original=strategies.sampled_from(random_input_data()),
156 level=strategies.integers(min_value=1, max_value=5),
164 level=strategies.integers(min_value=1, max_value=5),
157 source_read_size=strategies.integers(1, 16384),
165 source_read_size=strategies.integers(1, 16384),
158 read_size=strategies.integers(1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
166 read_size=strategies.integers(
167 1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
168 ),
159 )
169 )
160 def test_stream_source_readinto(self, original, level, source_read_size, read_size):
170 def test_stream_source_readinto(
171 self, original, level, source_read_size, read_size
172 ):
161 refctx = zstd.ZstdCompressor(level=level)
173 refctx = zstd.ZstdCompressor(level=level)
162 ref_frame = refctx.compress(original)
174 ref_frame = refctx.compress(original)
163
175
@@ -184,9 +196,13 b' class TestCompressor_stream_reader_fuzzi'
184 original=strategies.sampled_from(random_input_data()),
196 original=strategies.sampled_from(random_input_data()),
185 level=strategies.integers(min_value=1, max_value=5),
197 level=strategies.integers(min_value=1, max_value=5),
186 source_read_size=strategies.integers(1, 16384),
198 source_read_size=strategies.integers(1, 16384),
187 read_size=strategies.integers(1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
199 read_size=strategies.integers(
200 1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
201 ),
188 )
202 )
189 def test_buffer_source_readinto(self, original, level, source_read_size, read_size):
203 def test_buffer_source_readinto(
204 self, original, level, source_read_size, read_size
205 ):
190
206
191 refctx = zstd.ZstdCompressor(level=level)
207 refctx = zstd.ZstdCompressor(level=level)
192 ref_frame = refctx.compress(original)
208 ref_frame = refctx.compress(original)
@@ -285,9 +301,13 b' class TestCompressor_stream_reader_fuzzi'
285 original=strategies.sampled_from(random_input_data()),
301 original=strategies.sampled_from(random_input_data()),
286 level=strategies.integers(min_value=1, max_value=5),
302 level=strategies.integers(min_value=1, max_value=5),
287 source_read_size=strategies.integers(1, 16384),
303 source_read_size=strategies.integers(1, 16384),
288 read_size=strategies.integers(-1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
304 read_size=strategies.integers(
305 -1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
306 ),
289 )
307 )
290 def test_stream_source_read1(self, original, level, source_read_size, read_size):
308 def test_stream_source_read1(
309 self, original, level, source_read_size, read_size
310 ):
291 if read_size == 0:
311 if read_size == 0:
292 read_size = -1
312 read_size = -1
293
313
@@ -315,9 +335,13 b' class TestCompressor_stream_reader_fuzzi'
315 original=strategies.sampled_from(random_input_data()),
335 original=strategies.sampled_from(random_input_data()),
316 level=strategies.integers(min_value=1, max_value=5),
336 level=strategies.integers(min_value=1, max_value=5),
317 source_read_size=strategies.integers(1, 16384),
337 source_read_size=strategies.integers(1, 16384),
318 read_size=strategies.integers(-1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
338 read_size=strategies.integers(
339 -1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
340 ),
319 )
341 )
320 def test_buffer_source_read1(self, original, level, source_read_size, read_size):
342 def test_buffer_source_read1(
343 self, original, level, source_read_size, read_size
344 ):
321 if read_size == 0:
345 if read_size == 0:
322 read_size = -1
346 read_size = -1
323
347
@@ -412,7 +436,9 b' class TestCompressor_stream_reader_fuzzi'
412 original=strategies.sampled_from(random_input_data()),
436 original=strategies.sampled_from(random_input_data()),
413 level=strategies.integers(min_value=1, max_value=5),
437 level=strategies.integers(min_value=1, max_value=5),
414 source_read_size=strategies.integers(1, 16384),
438 source_read_size=strategies.integers(1, 16384),
415 read_size=strategies.integers(1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
439 read_size=strategies.integers(
440 1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
441 ),
416 )
442 )
417 def test_stream_source_readinto1(
443 def test_stream_source_readinto1(
418 self, original, level, source_read_size, read_size
444 self, original, level, source_read_size, read_size
@@ -446,7 +472,9 b' class TestCompressor_stream_reader_fuzzi'
446 original=strategies.sampled_from(random_input_data()),
472 original=strategies.sampled_from(random_input_data()),
447 level=strategies.integers(min_value=1, max_value=5),
473 level=strategies.integers(min_value=1, max_value=5),
448 source_read_size=strategies.integers(1, 16384),
474 source_read_size=strategies.integers(1, 16384),
449 read_size=strategies.integers(1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE),
475 read_size=strategies.integers(
476 1, zstd.COMPRESSION_RECOMMENDED_OUTPUT_SIZE
477 ),
450 )
478 )
451 def test_buffer_source_readinto1(
479 def test_buffer_source_readinto1(
452 self, original, level, source_read_size, read_size
480 self, original, level, source_read_size, read_size
@@ -576,7 +604,9 b' class TestCompressor_copy_stream_fuzzing'
576 read_size=strategies.integers(min_value=1, max_value=1048576),
604 read_size=strategies.integers(min_value=1, max_value=1048576),
577 write_size=strategies.integers(min_value=1, max_value=1048576),
605 write_size=strategies.integers(min_value=1, max_value=1048576),
578 )
606 )
579 def test_read_write_size_variance(self, original, level, read_size, write_size):
607 def test_read_write_size_variance(
608 self, original, level, read_size, write_size
609 ):
580 refctx = zstd.ZstdCompressor(level=level)
610 refctx = zstd.ZstdCompressor(level=level)
581 ref_frame = refctx.compress(original)
611 ref_frame = refctx.compress(original)
582
612
@@ -585,7 +615,11 b' class TestCompressor_copy_stream_fuzzing'
585 dest = io.BytesIO()
615 dest = io.BytesIO()
586
616
587 cctx.copy_stream(
617 cctx.copy_stream(
588 source, dest, size=len(original), read_size=read_size, write_size=write_size
618 source,
619 dest,
620 size=len(original),
621 read_size=read_size,
622 write_size=write_size,
589 )
623 )
590
624
591 self.assertEqual(dest.getvalue(), ref_frame)
625 self.assertEqual(dest.getvalue(), ref_frame)
@@ -675,7 +709,9 b' class TestCompressor_compressobj_fuzzing'
675 decompressed_chunks.append(dobj.decompress(chunk))
709 decompressed_chunks.append(dobj.decompress(chunk))
676
710
677 self.assertEqual(
711 self.assertEqual(
678 dctx.decompress(b"".join(compressed_chunks), max_output_size=len(original)),
712 dctx.decompress(
713 b"".join(compressed_chunks), max_output_size=len(original)
714 ),
679 original,
715 original,
680 )
716 )
681 self.assertEqual(b"".join(decompressed_chunks), original)
717 self.assertEqual(b"".join(decompressed_chunks), original)
@@ -690,7 +726,9 b' class TestCompressor_read_to_iter_fuzzin'
690 read_size=strategies.integers(min_value=1, max_value=4096),
726 read_size=strategies.integers(min_value=1, max_value=4096),
691 write_size=strategies.integers(min_value=1, max_value=4096),
727 write_size=strategies.integers(min_value=1, max_value=4096),
692 )
728 )
693 def test_read_write_size_variance(self, original, level, read_size, write_size):
729 def test_read_write_size_variance(
730 self, original, level, read_size, write_size
731 ):
694 refcctx = zstd.ZstdCompressor(level=level)
732 refcctx = zstd.ZstdCompressor(level=level)
695 ref_frame = refcctx.compress(original)
733 ref_frame = refcctx.compress(original)
696
734
@@ -699,7 +737,10 b' class TestCompressor_read_to_iter_fuzzin'
699 cctx = zstd.ZstdCompressor(level=level)
737 cctx = zstd.ZstdCompressor(level=level)
700 chunks = list(
738 chunks = list(
701 cctx.read_to_iter(
739 cctx.read_to_iter(
702 source, size=len(original), read_size=read_size, write_size=write_size
740 source,
741 size=len(original),
742 read_size=read_size,
743 write_size=write_size,
703 )
744 )
704 )
745 )
705
746
@@ -710,7 +751,9 b' class TestCompressor_read_to_iter_fuzzin'
710 class TestCompressor_multi_compress_to_buffer_fuzzing(TestCase):
751 class TestCompressor_multi_compress_to_buffer_fuzzing(TestCase):
711 @hypothesis.given(
752 @hypothesis.given(
712 original=strategies.lists(
753 original=strategies.lists(
713 strategies.sampled_from(random_input_data()), min_size=1, max_size=1024
754 strategies.sampled_from(random_input_data()),
755 min_size=1,
756 max_size=1024,
714 ),
757 ),
715 threads=strategies.integers(min_value=1, max_value=8),
758 threads=strategies.integers(min_value=1, max_value=8),
716 use_dict=strategies.booleans(),
759 use_dict=strategies.booleans(),
@@ -776,7 +819,8 b' class TestCompressor_chunker_fuzzing(Tes'
776 dctx = zstd.ZstdDecompressor()
819 dctx = zstd.ZstdDecompressor()
777
820
778 self.assertEqual(
821 self.assertEqual(
779 dctx.decompress(b"".join(chunks), max_output_size=len(original)), original
822 dctx.decompress(b"".join(chunks), max_output_size=len(original)),
823 original,
780 )
824 )
781
825
782 self.assertTrue(all(len(chunk) == chunk_size for chunk in chunks[:-1]))
826 self.assertTrue(all(len(chunk) == chunk_size for chunk in chunks[:-1]))
@@ -794,7 +838,9 b' class TestCompressor_chunker_fuzzing(Tes'
794 input_sizes=strategies.data(),
838 input_sizes=strategies.data(),
795 flushes=strategies.data(),
839 flushes=strategies.data(),
796 )
840 )
797 def test_flush_block(self, original, level, chunk_size, input_sizes, flushes):
841 def test_flush_block(
842 self, original, level, chunk_size, input_sizes, flushes
843 ):
798 cctx = zstd.ZstdCompressor(level=level)
844 cctx = zstd.ZstdCompressor(level=level)
799 chunker = cctx.chunker(chunk_size=chunk_size)
845 chunker = cctx.chunker(chunk_size=chunk_size)
800
846
@@ -830,7 +876,9 b' class TestCompressor_chunker_fuzzing(Tes'
830 decompressed_chunks.append(dobj.decompress(b"".join(chunks)))
876 decompressed_chunks.append(dobj.decompress(b"".join(chunks)))
831
877
832 self.assertEqual(
878 self.assertEqual(
833 dctx.decompress(b"".join(compressed_chunks), max_output_size=len(original)),
879 dctx.decompress(
880 b"".join(compressed_chunks), max_output_size=len(original)
881 ),
834 original,
882 original,
835 )
883 )
836 self.assertEqual(b"".join(decompressed_chunks), original)
884 self.assertEqual(b"".join(decompressed_chunks), original)
@@ -65,7 +65,9 b' class TestCompressionParameters(TestCase'
65 p = zstd.ZstdCompressionParameters(threads=4)
65 p = zstd.ZstdCompressionParameters(threads=4)
66 self.assertEqual(p.threads, 4)
66 self.assertEqual(p.threads, 4)
67
67
68 p = zstd.ZstdCompressionParameters(threads=2, job_size=1048576, overlap_log=6)
68 p = zstd.ZstdCompressionParameters(
69 threads=2, job_size=1048576, overlap_log=6
70 )
69 self.assertEqual(p.threads, 2)
71 self.assertEqual(p.threads, 2)
70 self.assertEqual(p.job_size, 1048576)
72 self.assertEqual(p.job_size, 1048576)
71 self.assertEqual(p.overlap_log, 6)
73 self.assertEqual(p.overlap_log, 6)
@@ -128,7 +130,9 b' class TestCompressionParameters(TestCase'
128 with self.assertRaisesRegex(
130 with self.assertRaisesRegex(
129 ValueError, "cannot specify both ldm_hash_rate_log"
131 ValueError, "cannot specify both ldm_hash_rate_log"
130 ):
132 ):
131 zstd.ZstdCompressionParameters(ldm_hash_rate_log=8, ldm_hash_every_log=4)
133 zstd.ZstdCompressionParameters(
134 ldm_hash_rate_log=8, ldm_hash_every_log=4
135 )
132
136
133 p = zstd.ZstdCompressionParameters(ldm_hash_rate_log=8)
137 p = zstd.ZstdCompressionParameters(ldm_hash_rate_log=8)
134 self.assertEqual(p.ldm_hash_every_log, 8)
138 self.assertEqual(p.ldm_hash_every_log, 8)
@@ -137,7 +141,9 b' class TestCompressionParameters(TestCase'
137 self.assertEqual(p.ldm_hash_every_log, 16)
141 self.assertEqual(p.ldm_hash_every_log, 16)
138
142
139 def test_overlap_log(self):
143 def test_overlap_log(self):
140 with self.assertRaisesRegex(ValueError, "cannot specify both overlap_log"):
144 with self.assertRaisesRegex(
145 ValueError, "cannot specify both overlap_log"
146 ):
141 zstd.ZstdCompressionParameters(overlap_log=1, overlap_size_log=9)
147 zstd.ZstdCompressionParameters(overlap_log=1, overlap_size_log=9)
142
148
143 p = zstd.ZstdCompressionParameters(overlap_log=2)
149 p = zstd.ZstdCompressionParameters(overlap_log=2)
@@ -169,10 +175,14 b' class TestFrameParameters(TestCase):'
169 zstd.get_frame_parameters(u"foobarbaz")
175 zstd.get_frame_parameters(u"foobarbaz")
170
176
171 def test_invalid_input_sizes(self):
177 def test_invalid_input_sizes(self):
172 with self.assertRaisesRegex(zstd.ZstdError, "not enough data for frame"):
178 with self.assertRaisesRegex(
179 zstd.ZstdError, "not enough data for frame"
180 ):
173 zstd.get_frame_parameters(b"")
181 zstd.get_frame_parameters(b"")
174
182
175 with self.assertRaisesRegex(zstd.ZstdError, "not enough data for frame"):
183 with self.assertRaisesRegex(
184 zstd.ZstdError, "not enough data for frame"
185 ):
176 zstd.get_frame_parameters(zstd.FRAME_HEADER)
186 zstd.get_frame_parameters(zstd.FRAME_HEADER)
177
187
178 def test_invalid_frame(self):
188 def test_invalid_frame(self):
@@ -201,7 +211,9 b' class TestFrameParameters(TestCase):'
201 self.assertTrue(params.has_checksum)
211 self.assertTrue(params.has_checksum)
202
212
203 # Upper 2 bits indicate content size.
213 # Upper 2 bits indicate content size.
204 params = zstd.get_frame_parameters(zstd.FRAME_HEADER + b"\x40\x00\xff\x00")
214 params = zstd.get_frame_parameters(
215 zstd.FRAME_HEADER + b"\x40\x00\xff\x00"
216 )
205 self.assertEqual(params.content_size, 511)
217 self.assertEqual(params.content_size, 511)
206 self.assertEqual(params.window_size, 1024)
218 self.assertEqual(params.window_size, 1024)
207 self.assertEqual(params.dict_id, 0)
219 self.assertEqual(params.dict_id, 0)
@@ -215,7 +227,9 b' class TestFrameParameters(TestCase):'
215 self.assertFalse(params.has_checksum)
227 self.assertFalse(params.has_checksum)
216
228
217 # Set multiple things.
229 # Set multiple things.
218 params = zstd.get_frame_parameters(zstd.FRAME_HEADER + b"\x45\x40\x0f\x10\x00")
230 params = zstd.get_frame_parameters(
231 zstd.FRAME_HEADER + b"\x45\x40\x0f\x10\x00"
232 )
219 self.assertEqual(params.content_size, 272)
233 self.assertEqual(params.content_size, 272)
220 self.assertEqual(params.window_size, 262144)
234 self.assertEqual(params.window_size, 262144)
221 self.assertEqual(params.dict_id, 15)
235 self.assertEqual(params.dict_id, 15)
@@ -23,7 +23,9 b' s_windowlog = strategies.integers('
23 s_chainlog = strategies.integers(
23 s_chainlog = strategies.integers(
24 min_value=zstd.CHAINLOG_MIN, max_value=zstd.CHAINLOG_MAX
24 min_value=zstd.CHAINLOG_MIN, max_value=zstd.CHAINLOG_MAX
25 )
25 )
26 s_hashlog = strategies.integers(min_value=zstd.HASHLOG_MIN, max_value=zstd.HASHLOG_MAX)
26 s_hashlog = strategies.integers(
27 min_value=zstd.HASHLOG_MIN, max_value=zstd.HASHLOG_MAX
28 )
27 s_searchlog = strategies.integers(
29 s_searchlog = strategies.integers(
28 min_value=zstd.SEARCHLOG_MIN, max_value=zstd.SEARCHLOG_MAX
30 min_value=zstd.SEARCHLOG_MIN, max_value=zstd.SEARCHLOG_MAX
29 )
31 )
@@ -61,7 +63,14 b' class TestCompressionParametersHypothesi'
61 s_strategy,
63 s_strategy,
62 )
64 )
63 def test_valid_init(
65 def test_valid_init(
64 self, windowlog, chainlog, hashlog, searchlog, minmatch, targetlength, strategy
66 self,
67 windowlog,
68 chainlog,
69 hashlog,
70 searchlog,
71 minmatch,
72 targetlength,
73 strategy,
65 ):
74 ):
66 zstd.ZstdCompressionParameters(
75 zstd.ZstdCompressionParameters(
67 window_log=windowlog,
76 window_log=windowlog,
@@ -83,7 +92,14 b' class TestCompressionParametersHypothesi'
83 s_strategy,
92 s_strategy,
84 )
93 )
85 def test_estimated_compression_context_size(
94 def test_estimated_compression_context_size(
86 self, windowlog, chainlog, hashlog, searchlog, minmatch, targetlength, strategy
95 self,
96 windowlog,
97 chainlog,
98 hashlog,
99 searchlog,
100 minmatch,
101 targetlength,
102 strategy,
87 ):
103 ):
88 if minmatch == zstd.MINMATCH_MIN and strategy in (
104 if minmatch == zstd.MINMATCH_MIN and strategy in (
89 zstd.STRATEGY_FAST,
105 zstd.STRATEGY_FAST,
@@ -170,11 +170,15 b' class TestDecompressor_decompress(TestCa'
170 dctx.decompress(compressed, max_output_size=len(source) - 1)
170 dctx.decompress(compressed, max_output_size=len(source) - 1)
171
171
172 # Input size + 1 works
172 # Input size + 1 works
173 decompressed = dctx.decompress(compressed, max_output_size=len(source) + 1)
173 decompressed = dctx.decompress(
174 compressed, max_output_size=len(source) + 1
175 )
174 self.assertEqual(decompressed, source)
176 self.assertEqual(decompressed, source)
175
177
176 # A much larger buffer works.
178 # A much larger buffer works.
177 decompressed = dctx.decompress(compressed, max_output_size=len(source) * 64)
179 decompressed = dctx.decompress(
180 compressed, max_output_size=len(source) * 64
181 )
178 self.assertEqual(decompressed, source)
182 self.assertEqual(decompressed, source)
179
183
180 def test_stupidly_large_output_buffer(self):
184 def test_stupidly_large_output_buffer(self):
@@ -237,7 +241,8 b' class TestDecompressor_decompress(TestCa'
237 dctx = zstd.ZstdDecompressor(max_window_size=2 ** zstd.WINDOWLOG_MIN)
241 dctx = zstd.ZstdDecompressor(max_window_size=2 ** zstd.WINDOWLOG_MIN)
238
242
239 with self.assertRaisesRegex(
243 with self.assertRaisesRegex(
240 zstd.ZstdError, "decompression error: Frame requires too much memory"
244 zstd.ZstdError,
245 "decompression error: Frame requires too much memory",
241 ):
246 ):
242 dctx.decompress(frame, max_output_size=len(source))
247 dctx.decompress(frame, max_output_size=len(source))
243
248
@@ -291,7 +296,9 b' class TestDecompressor_copy_stream(TestC'
291 self.assertEqual(w, len(source.getvalue()))
296 self.assertEqual(w, len(source.getvalue()))
292
297
293 def test_read_write_size(self):
298 def test_read_write_size(self):
294 source = OpCountingBytesIO(zstd.ZstdCompressor().compress(b"foobarfoobar"))
299 source = OpCountingBytesIO(
300 zstd.ZstdCompressor().compress(b"foobarfoobar")
301 )
295
302
296 dest = OpCountingBytesIO()
303 dest = OpCountingBytesIO()
297 dctx = zstd.ZstdDecompressor()
304 dctx = zstd.ZstdDecompressor()
@@ -309,7 +316,9 b' class TestDecompressor_stream_reader(Tes'
309 dctx = zstd.ZstdDecompressor()
316 dctx = zstd.ZstdDecompressor()
310
317
311 with dctx.stream_reader(b"foo") as reader:
318 with dctx.stream_reader(b"foo") as reader:
312 with self.assertRaisesRegex(ValueError, "cannot __enter__ multiple times"):
319 with self.assertRaisesRegex(
320 ValueError, "cannot __enter__ multiple times"
321 ):
313 with reader as reader2:
322 with reader as reader2:
314 pass
323 pass
315
324
@@ -474,7 +483,9 b' class TestDecompressor_stream_reader(Tes'
474 dctx = zstd.ZstdDecompressor()
483 dctx = zstd.ZstdDecompressor()
475
484
476 with dctx.stream_reader(frame) as reader:
485 with dctx.stream_reader(frame) as reader:
477 with self.assertRaisesRegex(ValueError, "cannot seek to negative position"):
486 with self.assertRaisesRegex(
487 ValueError, "cannot seek to negative position"
488 ):
478 reader.seek(-1, os.SEEK_SET)
489 reader.seek(-1, os.SEEK_SET)
479
490
480 reader.read(1)
491 reader.read(1)
@@ -490,7 +501,8 b' class TestDecompressor_stream_reader(Tes'
490 reader.seek(-1, os.SEEK_CUR)
501 reader.seek(-1, os.SEEK_CUR)
491
502
492 with self.assertRaisesRegex(
503 with self.assertRaisesRegex(
493 ValueError, "zstd decompression streams cannot be seeked with SEEK_END"
504 ValueError,
505 "zstd decompression streams cannot be seeked with SEEK_END",
494 ):
506 ):
495 reader.seek(0, os.SEEK_END)
507 reader.seek(0, os.SEEK_END)
496
508
@@ -743,7 +755,9 b' class TestDecompressor_stream_reader(Tes'
743
755
744 def test_read_lines(self):
756 def test_read_lines(self):
745 cctx = zstd.ZstdCompressor()
757 cctx = zstd.ZstdCompressor()
746 source = b"\n".join(("line %d" % i).encode("ascii") for i in range(1024))
758 source = b"\n".join(
759 ("line %d" % i).encode("ascii") for i in range(1024)
760 )
747
761
748 frame = cctx.compress(source)
762 frame = cctx.compress(source)
749
763
@@ -821,7 +835,9 b' class TestDecompressor_decompressobj(Tes'
821 dobj = dctx.decompressobj()
835 dobj = dctx.decompressobj()
822 dobj.decompress(data)
836 dobj.decompress(data)
823
837
824 with self.assertRaisesRegex(zstd.ZstdError, "cannot use a decompressobj"):
838 with self.assertRaisesRegex(
839 zstd.ZstdError, "cannot use a decompressobj"
840 ):
825 dobj.decompress(data)
841 dobj.decompress(data)
826 self.assertIsNone(dobj.flush())
842 self.assertIsNone(dobj.flush())
827
843
@@ -1124,7 +1140,9 b' class TestDecompressor_read_to_iter(Test'
1124 # Buffer protocol works.
1140 # Buffer protocol works.
1125 dctx.read_to_iter(b"foobar")
1141 dctx.read_to_iter(b"foobar")
1126
1142
1127 with self.assertRaisesRegex(ValueError, "must pass an object with a read"):
1143 with self.assertRaisesRegex(
1144 ValueError, "must pass an object with a read"
1145 ):
1128 b"".join(dctx.read_to_iter(True))
1146 b"".join(dctx.read_to_iter(True))
1129
1147
1130 def test_empty_input(self):
1148 def test_empty_input(self):
@@ -1226,7 +1244,9 b' class TestDecompressor_read_to_iter(Test'
1226 decompressed = b"".join(chunks)
1244 decompressed = b"".join(chunks)
1227 self.assertEqual(decompressed, source.getvalue())
1245 self.assertEqual(decompressed, source.getvalue())
1228
1246
1229 @unittest.skipUnless("ZSTD_SLOW_TESTS" in os.environ, "ZSTD_SLOW_TESTS not set")
1247 @unittest.skipUnless(
1248 "ZSTD_SLOW_TESTS" in os.environ, "ZSTD_SLOW_TESTS not set"
1249 )
1230 def test_large_input(self):
1250 def test_large_input(self):
1231 bytes = list(struct.Struct(">B").pack(i) for i in range(256))
1251 bytes = list(struct.Struct(">B").pack(i) for i in range(256))
1232 compressed = NonClosingBytesIO()
1252 compressed = NonClosingBytesIO()
@@ -1241,13 +1261,16 b' class TestDecompressor_read_to_iter(Test'
1241 len(compressed.getvalue())
1261 len(compressed.getvalue())
1242 > zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE
1262 > zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE
1243 )
1263 )
1244 have_raw = input_size > zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE * 2
1264 have_raw = (
1265 input_size > zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE * 2
1266 )
1245 if have_compressed and have_raw:
1267 if have_compressed and have_raw:
1246 break
1268 break
1247
1269
1248 compressed = io.BytesIO(compressed.getvalue())
1270 compressed = io.BytesIO(compressed.getvalue())
1249 self.assertGreater(
1271 self.assertGreater(
1250 len(compressed.getvalue()), zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE
1272 len(compressed.getvalue()),
1273 zstd.DECOMPRESSION_RECOMMENDED_INPUT_SIZE,
1251 )
1274 )
1252
1275
1253 dctx = zstd.ZstdDecompressor()
1276 dctx = zstd.ZstdDecompressor()
@@ -1303,7 +1326,9 b' class TestDecompressor_read_to_iter(Test'
1303 self.assertEqual(streamed, source.getvalue())
1326 self.assertEqual(streamed, source.getvalue())
1304
1327
1305 def test_read_write_size(self):
1328 def test_read_write_size(self):
1306 source = OpCountingBytesIO(zstd.ZstdCompressor().compress(b"foobarfoobar"))
1329 source = OpCountingBytesIO(
1330 zstd.ZstdCompressor().compress(b"foobarfoobar")
1331 )
1307 dctx = zstd.ZstdDecompressor()
1332 dctx = zstd.ZstdDecompressor()
1308 for chunk in dctx.read_to_iter(source, read_size=1, write_size=1):
1333 for chunk in dctx.read_to_iter(source, read_size=1, write_size=1):
1309 self.assertEqual(len(chunk), 1)
1334 self.assertEqual(len(chunk), 1)
@@ -1355,10 +1380,14 b' class TestDecompressor_content_dict_chai'
1355 ):
1380 ):
1356 dctx.decompress_content_dict_chain([zstd.FRAME_HEADER])
1381 dctx.decompress_content_dict_chain([zstd.FRAME_HEADER])
1357
1382
1358 with self.assertRaisesRegex(ValueError, "chunk 0 is not a valid zstd frame"):
1383 with self.assertRaisesRegex(
1384 ValueError, "chunk 0 is not a valid zstd frame"
1385 ):
1359 dctx.decompress_content_dict_chain([b"foo" * 8])
1386 dctx.decompress_content_dict_chain([b"foo" * 8])
1360
1387
1361 no_size = zstd.ZstdCompressor(write_content_size=False).compress(b"foo" * 64)
1388 no_size = zstd.ZstdCompressor(write_content_size=False).compress(
1389 b"foo" * 64
1390 )
1362
1391
1363 with self.assertRaisesRegex(
1392 with self.assertRaisesRegex(
1364 ValueError, "chunk 0 missing content size in frame"
1393 ValueError, "chunk 0 missing content size in frame"
@@ -1389,10 +1418,14 b' class TestDecompressor_content_dict_chai'
1389 ):
1418 ):
1390 dctx.decompress_content_dict_chain([initial, zstd.FRAME_HEADER])
1419 dctx.decompress_content_dict_chain([initial, zstd.FRAME_HEADER])
1391
1420
1392 with self.assertRaisesRegex(ValueError, "chunk 1 is not a valid zstd frame"):
1421 with self.assertRaisesRegex(
1422 ValueError, "chunk 1 is not a valid zstd frame"
1423 ):
1393 dctx.decompress_content_dict_chain([initial, b"foo" * 8])
1424 dctx.decompress_content_dict_chain([initial, b"foo" * 8])
1394
1425
1395 no_size = zstd.ZstdCompressor(write_content_size=False).compress(b"foo" * 64)
1426 no_size = zstd.ZstdCompressor(write_content_size=False).compress(
1427 b"foo" * 64
1428 )
1396
1429
1397 with self.assertRaisesRegex(
1430 with self.assertRaisesRegex(
1398 ValueError, "chunk 1 missing content size in frame"
1431 ValueError, "chunk 1 missing content size in frame"
@@ -1400,7 +1433,9 b' class TestDecompressor_content_dict_chai'
1400 dctx.decompress_content_dict_chain([initial, no_size])
1433 dctx.decompress_content_dict_chain([initial, no_size])
1401
1434
1402 # Corrupt second frame.
1435 # Corrupt second frame.
1403 cctx = zstd.ZstdCompressor(dict_data=zstd.ZstdCompressionDict(b"foo" * 64))
1436 cctx = zstd.ZstdCompressor(
1437 dict_data=zstd.ZstdCompressionDict(b"foo" * 64)
1438 )
1404 frame = cctx.compress(b"bar" * 64)
1439 frame = cctx.compress(b"bar" * 64)
1405 frame = frame[0:12] + frame[15:]
1440 frame = frame[0:12] + frame[15:]
1406
1441
@@ -1447,7 +1482,9 b' class TestDecompressor_multi_decompress_'
1447 with self.assertRaises(TypeError):
1482 with self.assertRaises(TypeError):
1448 dctx.multi_decompress_to_buffer((1, 2))
1483 dctx.multi_decompress_to_buffer((1, 2))
1449
1484
1450 with self.assertRaisesRegex(TypeError, "item 0 not a bytes like object"):
1485 with self.assertRaisesRegex(
1486 TypeError, "item 0 not a bytes like object"
1487 ):
1451 dctx.multi_decompress_to_buffer([u"foo"])
1488 dctx.multi_decompress_to_buffer([u"foo"])
1452
1489
1453 with self.assertRaisesRegex(
1490 with self.assertRaisesRegex(
@@ -1491,7 +1528,9 b' class TestDecompressor_multi_decompress_'
1491 if not hasattr(dctx, "multi_decompress_to_buffer"):
1528 if not hasattr(dctx, "multi_decompress_to_buffer"):
1492 self.skipTest("multi_decompress_to_buffer not available")
1529 self.skipTest("multi_decompress_to_buffer not available")
1493
1530
1494 result = dctx.multi_decompress_to_buffer(frames, decompressed_sizes=sizes)
1531 result = dctx.multi_decompress_to_buffer(
1532 frames, decompressed_sizes=sizes
1533 )
1495
1534
1496 self.assertEqual(len(result), len(frames))
1535 self.assertEqual(len(result), len(frames))
1497 self.assertEqual(result.size(), sum(map(len, original)))
1536 self.assertEqual(result.size(), sum(map(len, original)))
@@ -1582,10 +1621,15 b' class TestDecompressor_multi_decompress_'
1582 # And a manual mode.
1621 # And a manual mode.
1583 b = b"".join([frames[0].tobytes(), frames[1].tobytes()])
1622 b = b"".join([frames[0].tobytes(), frames[1].tobytes()])
1584 b1 = zstd.BufferWithSegments(
1623 b1 = zstd.BufferWithSegments(
1585 b, struct.pack("=QQQQ", 0, len(frames[0]), len(frames[0]), len(frames[1]))
1624 b,
1625 struct.pack(
1626 "=QQQQ", 0, len(frames[0]), len(frames[0]), len(frames[1])
1627 ),
1586 )
1628 )
1587
1629
1588 b = b"".join([frames[2].tobytes(), frames[3].tobytes(), frames[4].tobytes()])
1630 b = b"".join(
1631 [frames[2].tobytes(), frames[3].tobytes(), frames[4].tobytes()]
1632 )
1589 b2 = zstd.BufferWithSegments(
1633 b2 = zstd.BufferWithSegments(
1590 b,
1634 b,
1591 struct.pack(
1635 struct.pack(
@@ -196,7 +196,9 b' class TestDecompressor_stream_reader_fuz'
196 streaming=strategies.booleans(),
196 streaming=strategies.booleans(),
197 source_read_size=strategies.integers(1, 1048576),
197 source_read_size=strategies.integers(1, 1048576),
198 )
198 )
199 def test_stream_source_readall(self, original, level, streaming, source_read_size):
199 def test_stream_source_readall(
200 self, original, level, streaming, source_read_size
201 ):
200 cctx = zstd.ZstdCompressor(level=level)
202 cctx = zstd.ZstdCompressor(level=level)
201
203
202 if streaming:
204 if streaming:
@@ -398,7 +400,9 b' class TestDecompressor_stream_writer_fuz'
398 write_size=strategies.integers(min_value=1, max_value=8192),
400 write_size=strategies.integers(min_value=1, max_value=8192),
399 input_sizes=strategies.data(),
401 input_sizes=strategies.data(),
400 )
402 )
401 def test_write_size_variance(self, original, level, write_size, input_sizes):
403 def test_write_size_variance(
404 self, original, level, write_size, input_sizes
405 ):
402 cctx = zstd.ZstdCompressor(level=level)
406 cctx = zstd.ZstdCompressor(level=level)
403 frame = cctx.compress(original)
407 frame = cctx.compress(original)
404
408
@@ -433,7 +437,9 b' class TestDecompressor_copy_stream_fuzzi'
433 read_size=strategies.integers(min_value=1, max_value=8192),
437 read_size=strategies.integers(min_value=1, max_value=8192),
434 write_size=strategies.integers(min_value=1, max_value=8192),
438 write_size=strategies.integers(min_value=1, max_value=8192),
435 )
439 )
436 def test_read_write_size_variance(self, original, level, read_size, write_size):
440 def test_read_write_size_variance(
441 self, original, level, read_size, write_size
442 ):
437 cctx = zstd.ZstdCompressor(level=level)
443 cctx = zstd.ZstdCompressor(level=level)
438 frame = cctx.compress(original)
444 frame = cctx.compress(original)
439
445
@@ -441,7 +447,9 b' class TestDecompressor_copy_stream_fuzzi'
441 dest = io.BytesIO()
447 dest = io.BytesIO()
442
448
443 dctx = zstd.ZstdDecompressor()
449 dctx = zstd.ZstdDecompressor()
444 dctx.copy_stream(source, dest, read_size=read_size, write_size=write_size)
450 dctx.copy_stream(
451 source, dest, read_size=read_size, write_size=write_size
452 )
445
453
446 self.assertEqual(dest.getvalue(), original)
454 self.assertEqual(dest.getvalue(), original)
447
455
@@ -490,11 +498,14 b' class TestDecompressor_decompressobj_fuz'
490 original=strategies.sampled_from(random_input_data()),
498 original=strategies.sampled_from(random_input_data()),
491 level=strategies.integers(min_value=1, max_value=5),
499 level=strategies.integers(min_value=1, max_value=5),
492 write_size=strategies.integers(
500 write_size=strategies.integers(
493 min_value=1, max_value=4 * zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE
501 min_value=1,
502 max_value=4 * zstd.DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE,
494 ),
503 ),
495 chunk_sizes=strategies.data(),
504 chunk_sizes=strategies.data(),
496 )
505 )
497 def test_random_output_sizes(self, original, level, write_size, chunk_sizes):
506 def test_random_output_sizes(
507 self, original, level, write_size, chunk_sizes
508 ):
498 cctx = zstd.ZstdCompressor(level=level)
509 cctx = zstd.ZstdCompressor(level=level)
499 frame = cctx.compress(original)
510 frame = cctx.compress(original)
500
511
@@ -524,7 +535,9 b' class TestDecompressor_read_to_iter_fuzz'
524 read_size=strategies.integers(min_value=1, max_value=4096),
535 read_size=strategies.integers(min_value=1, max_value=4096),
525 write_size=strategies.integers(min_value=1, max_value=4096),
536 write_size=strategies.integers(min_value=1, max_value=4096),
526 )
537 )
527 def test_read_write_size_variance(self, original, level, read_size, write_size):
538 def test_read_write_size_variance(
539 self, original, level, read_size, write_size
540 ):
528 cctx = zstd.ZstdCompressor(level=level)
541 cctx = zstd.ZstdCompressor(level=level)
529 frame = cctx.compress(original)
542 frame = cctx.compress(original)
530
543
@@ -532,7 +545,9 b' class TestDecompressor_read_to_iter_fuzz'
532
545
533 dctx = zstd.ZstdDecompressor()
546 dctx = zstd.ZstdDecompressor()
534 chunks = list(
547 chunks = list(
535 dctx.read_to_iter(source, read_size=read_size, write_size=write_size)
548 dctx.read_to_iter(
549 source, read_size=read_size, write_size=write_size
550 )
536 )
551 )
537
552
538 self.assertEqual(b"".join(chunks), original)
553 self.assertEqual(b"".join(chunks), original)
@@ -542,7 +557,9 b' class TestDecompressor_read_to_iter_fuzz'
542 class TestDecompressor_multi_decompress_to_buffer_fuzzing(TestCase):
557 class TestDecompressor_multi_decompress_to_buffer_fuzzing(TestCase):
543 @hypothesis.given(
558 @hypothesis.given(
544 original=strategies.lists(
559 original=strategies.lists(
545 strategies.sampled_from(random_input_data()), min_size=1, max_size=1024
560 strategies.sampled_from(random_input_data()),
561 min_size=1,
562 max_size=1024,
546 ),
563 ),
547 threads=strategies.integers(min_value=1, max_value=8),
564 threads=strategies.integers(min_value=1, max_value=8),
548 use_dict=strategies.booleans(),
565 use_dict=strategies.booleans(),
@@ -51,11 +51,15 b' class TestTrainDictionary(TestCase):'
51 self.assertEqual(d.d, 16)
51 self.assertEqual(d.d, 16)
52
52
53 def test_set_dict_id(self):
53 def test_set_dict_id(self):
54 d = zstd.train_dictionary(8192, generate_samples(), k=64, d=16, dict_id=42)
54 d = zstd.train_dictionary(
55 8192, generate_samples(), k=64, d=16, dict_id=42
56 )
55 self.assertEqual(d.dict_id(), 42)
57 self.assertEqual(d.dict_id(), 42)
56
58
57 def test_optimize(self):
59 def test_optimize(self):
58 d = zstd.train_dictionary(8192, generate_samples(), threads=-1, steps=1, d=16)
60 d = zstd.train_dictionary(
61 8192, generate_samples(), threads=-1, steps=1, d=16
62 )
59
63
60 # This varies by platform.
64 # This varies by platform.
61 self.assertIn(d.k, (50, 2000))
65 self.assertIn(d.k, (50, 2000))
@@ -71,10 +75,14 b' class TestCompressionDict(TestCase):'
71 def test_bad_precompute_compress(self):
75 def test_bad_precompute_compress(self):
72 d = zstd.train_dictionary(8192, generate_samples(), k=64, d=16)
76 d = zstd.train_dictionary(8192, generate_samples(), k=64, d=16)
73
77
74 with self.assertRaisesRegex(ValueError, "must specify one of level or "):
78 with self.assertRaisesRegex(
79 ValueError, "must specify one of level or "
80 ):
75 d.precompute_compress()
81 d.precompute_compress()
76
82
77 with self.assertRaisesRegex(ValueError, "must only specify one of level or "):
83 with self.assertRaisesRegex(
84 ValueError, "must only specify one of level or "
85 ):
78 d.precompute_compress(
86 d.precompute_compress(
79 level=3, compression_params=zstd.CompressionParameters()
87 level=3, compression_params=zstd.CompressionParameters()
80 )
88 )
@@ -88,5 +96,7 b' class TestCompressionDict(TestCase):'
88 d = zstd.ZstdCompressionDict(
96 d = zstd.ZstdCompressionDict(
89 b"dictcontent" * 64, dict_type=zstd.DICT_TYPE_FULLDICT
97 b"dictcontent" * 64, dict_type=zstd.DICT_TYPE_FULLDICT
90 )
98 )
91 with self.assertRaisesRegex(zstd.ZstdError, "unable to precompute dictionary"):
99 with self.assertRaisesRegex(
100 zstd.ZstdError, "unable to precompute dictionary"
101 ):
92 d.precompute_compress(level=1)
102 d.precompute_compress(level=1)
@@ -299,10 +299,14 b' class ZstdCompressionParameters(object):'
299 _set_compression_parameter(params, lib.ZSTD_c_chainLog, chain_log)
299 _set_compression_parameter(params, lib.ZSTD_c_chainLog, chain_log)
300 _set_compression_parameter(params, lib.ZSTD_c_searchLog, search_log)
300 _set_compression_parameter(params, lib.ZSTD_c_searchLog, search_log)
301 _set_compression_parameter(params, lib.ZSTD_c_minMatch, min_match)
301 _set_compression_parameter(params, lib.ZSTD_c_minMatch, min_match)
302 _set_compression_parameter(params, lib.ZSTD_c_targetLength, target_length)
302 _set_compression_parameter(
303 params, lib.ZSTD_c_targetLength, target_length
304 )
303
305
304 if strategy != -1 and compression_strategy != -1:
306 if strategy != -1 and compression_strategy != -1:
305 raise ValueError("cannot specify both compression_strategy and strategy")
307 raise ValueError(
308 "cannot specify both compression_strategy and strategy"
309 )
306
310
307 if compression_strategy != -1:
311 if compression_strategy != -1:
308 strategy = compression_strategy
312 strategy = compression_strategy
@@ -313,12 +317,16 b' class ZstdCompressionParameters(object):'
313 _set_compression_parameter(
317 _set_compression_parameter(
314 params, lib.ZSTD_c_contentSizeFlag, write_content_size
318 params, lib.ZSTD_c_contentSizeFlag, write_content_size
315 )
319 )
316 _set_compression_parameter(params, lib.ZSTD_c_checksumFlag, write_checksum)
320 _set_compression_parameter(
321 params, lib.ZSTD_c_checksumFlag, write_checksum
322 )
317 _set_compression_parameter(params, lib.ZSTD_c_dictIDFlag, write_dict_id)
323 _set_compression_parameter(params, lib.ZSTD_c_dictIDFlag, write_dict_id)
318 _set_compression_parameter(params, lib.ZSTD_c_jobSize, job_size)
324 _set_compression_parameter(params, lib.ZSTD_c_jobSize, job_size)
319
325
320 if overlap_log != -1 and overlap_size_log != -1:
326 if overlap_log != -1 and overlap_size_log != -1:
321 raise ValueError("cannot specify both overlap_log and overlap_size_log")
327 raise ValueError(
328 "cannot specify both overlap_log and overlap_size_log"
329 )
322
330
323 if overlap_size_log != -1:
331 if overlap_size_log != -1:
324 overlap_log = overlap_size_log
332 overlap_log = overlap_size_log
@@ -326,12 +334,16 b' class ZstdCompressionParameters(object):'
326 overlap_log = 0
334 overlap_log = 0
327
335
328 _set_compression_parameter(params, lib.ZSTD_c_overlapLog, overlap_log)
336 _set_compression_parameter(params, lib.ZSTD_c_overlapLog, overlap_log)
329 _set_compression_parameter(params, lib.ZSTD_c_forceMaxWindow, force_max_window)
337 _set_compression_parameter(
338 params, lib.ZSTD_c_forceMaxWindow, force_max_window
339 )
330 _set_compression_parameter(
340 _set_compression_parameter(
331 params, lib.ZSTD_c_enableLongDistanceMatching, enable_ldm
341 params, lib.ZSTD_c_enableLongDistanceMatching, enable_ldm
332 )
342 )
333 _set_compression_parameter(params, lib.ZSTD_c_ldmHashLog, ldm_hash_log)
343 _set_compression_parameter(params, lib.ZSTD_c_ldmHashLog, ldm_hash_log)
334 _set_compression_parameter(params, lib.ZSTD_c_ldmMinMatch, ldm_min_match)
344 _set_compression_parameter(
345 params, lib.ZSTD_c_ldmMinMatch, ldm_min_match
346 )
335 _set_compression_parameter(
347 _set_compression_parameter(
336 params, lib.ZSTD_c_ldmBucketSizeLog, ldm_bucket_size_log
348 params, lib.ZSTD_c_ldmBucketSizeLog, ldm_bucket_size_log
337 )
349 )
@@ -346,7 +358,9 b' class ZstdCompressionParameters(object):'
346 elif ldm_hash_rate_log == -1:
358 elif ldm_hash_rate_log == -1:
347 ldm_hash_rate_log = 0
359 ldm_hash_rate_log = 0
348
360
349 _set_compression_parameter(params, lib.ZSTD_c_ldmHashRateLog, ldm_hash_rate_log)
361 _set_compression_parameter(
362 params, lib.ZSTD_c_ldmHashRateLog, ldm_hash_rate_log
363 )
350
364
351 @property
365 @property
352 def format(self):
366 def format(self):
@@ -354,7 +368,9 b' class ZstdCompressionParameters(object):'
354
368
355 @property
369 @property
356 def compression_level(self):
370 def compression_level(self):
357 return _get_compression_parameter(self._params, lib.ZSTD_c_compressionLevel)
371 return _get_compression_parameter(
372 self._params, lib.ZSTD_c_compressionLevel
373 )
358
374
359 @property
375 @property
360 def window_log(self):
376 def window_log(self):
@@ -386,7 +402,9 b' class ZstdCompressionParameters(object):'
386
402
387 @property
403 @property
388 def write_content_size(self):
404 def write_content_size(self):
389 return _get_compression_parameter(self._params, lib.ZSTD_c_contentSizeFlag)
405 return _get_compression_parameter(
406 self._params, lib.ZSTD_c_contentSizeFlag
407 )
390
408
391 @property
409 @property
392 def write_checksum(self):
410 def write_checksum(self):
@@ -410,7 +428,9 b' class ZstdCompressionParameters(object):'
410
428
411 @property
429 @property
412 def force_max_window(self):
430 def force_max_window(self):
413 return _get_compression_parameter(self._params, lib.ZSTD_c_forceMaxWindow)
431 return _get_compression_parameter(
432 self._params, lib.ZSTD_c_forceMaxWindow
433 )
414
434
415 @property
435 @property
416 def enable_ldm(self):
436 def enable_ldm(self):
@@ -428,11 +448,15 b' class ZstdCompressionParameters(object):'
428
448
429 @property
449 @property
430 def ldm_bucket_size_log(self):
450 def ldm_bucket_size_log(self):
431 return _get_compression_parameter(self._params, lib.ZSTD_c_ldmBucketSizeLog)
451 return _get_compression_parameter(
452 self._params, lib.ZSTD_c_ldmBucketSizeLog
453 )
432
454
433 @property
455 @property
434 def ldm_hash_rate_log(self):
456 def ldm_hash_rate_log(self):
435 return _get_compression_parameter(self._params, lib.ZSTD_c_ldmHashRateLog)
457 return _get_compression_parameter(
458 self._params, lib.ZSTD_c_ldmHashRateLog
459 )
436
460
437 @property
461 @property
438 def ldm_hash_every_log(self):
462 def ldm_hash_every_log(self):
@@ -457,7 +481,8 b' def _set_compression_parameter(params, p'
457 zresult = lib.ZSTD_CCtxParams_setParameter(params, param, value)
481 zresult = lib.ZSTD_CCtxParams_setParameter(params, param, value)
458 if lib.ZSTD_isError(zresult):
482 if lib.ZSTD_isError(zresult):
459 raise ZstdError(
483 raise ZstdError(
460 "unable to set compression context parameter: %s" % _zstd_error(zresult)
484 "unable to set compression context parameter: %s"
485 % _zstd_error(zresult)
461 )
486 )
462
487
463
488
@@ -467,14 +492,17 b' def _get_compression_parameter(params, p'
467 zresult = lib.ZSTD_CCtxParams_getParameter(params, param, result)
492 zresult = lib.ZSTD_CCtxParams_getParameter(params, param, result)
468 if lib.ZSTD_isError(zresult):
493 if lib.ZSTD_isError(zresult):
469 raise ZstdError(
494 raise ZstdError(
470 "unable to get compression context parameter: %s" % _zstd_error(zresult)
495 "unable to get compression context parameter: %s"
496 % _zstd_error(zresult)
471 )
497 )
472
498
473 return result[0]
499 return result[0]
474
500
475
501
476 class ZstdCompressionWriter(object):
502 class ZstdCompressionWriter(object):
477 def __init__(self, compressor, writer, source_size, write_size, write_return_read):
503 def __init__(
504 self, compressor, writer, source_size, write_size, write_return_read
505 ):
478 self._compressor = compressor
506 self._compressor = compressor
479 self._writer = writer
507 self._writer = writer
480 self._write_size = write_size
508 self._write_size = write_size
@@ -491,7 +519,9 b' class ZstdCompressionWriter(object):'
491
519
492 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(compressor._cctx, source_size)
520 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(compressor._cctx, source_size)
493 if lib.ZSTD_isError(zresult):
521 if lib.ZSTD_isError(zresult):
494 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
522 raise ZstdError(
523 "error setting source size: %s" % _zstd_error(zresult)
524 )
495
525
496 def __enter__(self):
526 def __enter__(self):
497 if self._closed:
527 if self._closed:
@@ -595,13 +625,20 b' class ZstdCompressionWriter(object):'
595
625
596 while in_buffer.pos < in_buffer.size:
626 while in_buffer.pos < in_buffer.size:
597 zresult = lib.ZSTD_compressStream2(
627 zresult = lib.ZSTD_compressStream2(
598 self._compressor._cctx, out_buffer, in_buffer, lib.ZSTD_e_continue
628 self._compressor._cctx,
629 out_buffer,
630 in_buffer,
631 lib.ZSTD_e_continue,
599 )
632 )
600 if lib.ZSTD_isError(zresult):
633 if lib.ZSTD_isError(zresult):
601 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
634 raise ZstdError(
635 "zstd compress error: %s" % _zstd_error(zresult)
636 )
602
637
603 if out_buffer.pos:
638 if out_buffer.pos:
604 self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
639 self._writer.write(
640 ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
641 )
605 total_write += out_buffer.pos
642 total_write += out_buffer.pos
606 self._bytes_compressed += out_buffer.pos
643 self._bytes_compressed += out_buffer.pos
607 out_buffer.pos = 0
644 out_buffer.pos = 0
@@ -637,10 +674,14 b' class ZstdCompressionWriter(object):'
637 self._compressor._cctx, out_buffer, in_buffer, flush
674 self._compressor._cctx, out_buffer, in_buffer, flush
638 )
675 )
639 if lib.ZSTD_isError(zresult):
676 if lib.ZSTD_isError(zresult):
640 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
677 raise ZstdError(
678 "zstd compress error: %s" % _zstd_error(zresult)
679 )
641
680
642 if out_buffer.pos:
681 if out_buffer.pos:
643 self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
682 self._writer.write(
683 ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
684 )
644 total_write += out_buffer.pos
685 total_write += out_buffer.pos
645 self._bytes_compressed += out_buffer.pos
686 self._bytes_compressed += out_buffer.pos
646 out_buffer.pos = 0
687 out_buffer.pos = 0
@@ -672,7 +713,9 b' class ZstdCompressionObj(object):'
672 self._compressor._cctx, self._out, source, lib.ZSTD_e_continue
713 self._compressor._cctx, self._out, source, lib.ZSTD_e_continue
673 )
714 )
674 if lib.ZSTD_isError(zresult):
715 if lib.ZSTD_isError(zresult):
675 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
716 raise ZstdError(
717 "zstd compress error: %s" % _zstd_error(zresult)
718 )
676
719
677 if self._out.pos:
720 if self._out.pos:
678 chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])
721 chunks.append(ffi.buffer(self._out.dst, self._out.pos)[:])
@@ -681,7 +724,10 b' class ZstdCompressionObj(object):'
681 return b"".join(chunks)
724 return b"".join(chunks)
682
725
683 def flush(self, flush_mode=COMPRESSOBJ_FLUSH_FINISH):
726 def flush(self, flush_mode=COMPRESSOBJ_FLUSH_FINISH):
684 if flush_mode not in (COMPRESSOBJ_FLUSH_FINISH, COMPRESSOBJ_FLUSH_BLOCK):
727 if flush_mode not in (
728 COMPRESSOBJ_FLUSH_FINISH,
729 COMPRESSOBJ_FLUSH_BLOCK,
730 ):
685 raise ValueError("flush mode not recognized")
731 raise ValueError("flush mode not recognized")
686
732
687 if self._finished:
733 if self._finished:
@@ -768,7 +814,9 b' class ZstdCompressionChunker(object):'
768 self._in.pos = 0
814 self._in.pos = 0
769
815
770 if lib.ZSTD_isError(zresult):
816 if lib.ZSTD_isError(zresult):
771 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
817 raise ZstdError(
818 "zstd compress error: %s" % _zstd_error(zresult)
819 )
772
820
773 if self._out.pos == self._out.size:
821 if self._out.pos == self._out.size:
774 yield ffi.buffer(self._out.dst, self._out.pos)[:]
822 yield ffi.buffer(self._out.dst, self._out.pos)[:]
@@ -780,7 +828,8 b' class ZstdCompressionChunker(object):'
780
828
781 if self._in.src != ffi.NULL:
829 if self._in.src != ffi.NULL:
782 raise ZstdError(
830 raise ZstdError(
783 "cannot call flush() before consuming output from " "previous operation"
831 "cannot call flush() before consuming output from "
832 "previous operation"
784 )
833 )
785
834
786 while True:
835 while True:
@@ -788,7 +837,9 b' class ZstdCompressionChunker(object):'
788 self._compressor._cctx, self._out, self._in, lib.ZSTD_e_flush
837 self._compressor._cctx, self._out, self._in, lib.ZSTD_e_flush
789 )
838 )
790 if lib.ZSTD_isError(zresult):
839 if lib.ZSTD_isError(zresult):
791 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
840 raise ZstdError(
841 "zstd compress error: %s" % _zstd_error(zresult)
842 )
792
843
793 if self._out.pos:
844 if self._out.pos:
794 yield ffi.buffer(self._out.dst, self._out.pos)[:]
845 yield ffi.buffer(self._out.dst, self._out.pos)[:]
@@ -812,7 +863,9 b' class ZstdCompressionChunker(object):'
812 self._compressor._cctx, self._out, self._in, lib.ZSTD_e_end
863 self._compressor._cctx, self._out, self._in, lib.ZSTD_e_end
813 )
864 )
814 if lib.ZSTD_isError(zresult):
865 if lib.ZSTD_isError(zresult):
815 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
866 raise ZstdError(
867 "zstd compress error: %s" % _zstd_error(zresult)
868 )
816
869
817 if self._out.pos:
870 if self._out.pos:
818 yield ffi.buffer(self._out.dst, self._out.pos)[:]
871 yield ffi.buffer(self._out.dst, self._out.pos)[:]
@@ -939,7 +992,10 b' class ZstdCompressionReader(object):'
939 old_pos = out_buffer.pos
992 old_pos = out_buffer.pos
940
993
941 zresult = lib.ZSTD_compressStream2(
994 zresult = lib.ZSTD_compressStream2(
942 self._compressor._cctx, out_buffer, self._in_buffer, lib.ZSTD_e_continue
995 self._compressor._cctx,
996 out_buffer,
997 self._in_buffer,
998 lib.ZSTD_e_continue,
943 )
999 )
944
1000
945 self._bytes_compressed += out_buffer.pos - old_pos
1001 self._bytes_compressed += out_buffer.pos - old_pos
@@ -997,7 +1053,9 b' class ZstdCompressionReader(object):'
997 self._bytes_compressed += out_buffer.pos - old_pos
1053 self._bytes_compressed += out_buffer.pos - old_pos
998
1054
999 if lib.ZSTD_isError(zresult):
1055 if lib.ZSTD_isError(zresult):
1000 raise ZstdError("error ending compression stream: %s", _zstd_error(zresult))
1056 raise ZstdError(
1057 "error ending compression stream: %s", _zstd_error(zresult)
1058 )
1001
1059
1002 if zresult == 0:
1060 if zresult == 0:
1003 self._finished_output = True
1061 self._finished_output = True
@@ -1102,7 +1160,9 b' class ZstdCompressionReader(object):'
1102 self._bytes_compressed += out_buffer.pos - old_pos
1160 self._bytes_compressed += out_buffer.pos - old_pos
1103
1161
1104 if lib.ZSTD_isError(zresult):
1162 if lib.ZSTD_isError(zresult):
1105 raise ZstdError("error ending compression stream: %s", _zstd_error(zresult))
1163 raise ZstdError(
1164 "error ending compression stream: %s", _zstd_error(zresult)
1165 )
1106
1166
1107 if zresult == 0:
1167 if zresult == 0:
1108 self._finished_output = True
1168 self._finished_output = True
@@ -1170,13 +1230,17 b' class ZstdCompressor(object):'
1170 threads=0,
1230 threads=0,
1171 ):
1231 ):
1172 if level > lib.ZSTD_maxCLevel():
1232 if level > lib.ZSTD_maxCLevel():
1173 raise ValueError("level must be less than %d" % lib.ZSTD_maxCLevel())
1233 raise ValueError(
1234 "level must be less than %d" % lib.ZSTD_maxCLevel()
1235 )
1174
1236
1175 if threads < 0:
1237 if threads < 0:
1176 threads = _cpu_count()
1238 threads = _cpu_count()
1177
1239
1178 if compression_params and write_checksum is not None:
1240 if compression_params and write_checksum is not None:
1179 raise ValueError("cannot define compression_params and " "write_checksum")
1241 raise ValueError(
1242 "cannot define compression_params and " "write_checksum"
1243 )
1180
1244
1181 if compression_params and write_content_size is not None:
1245 if compression_params and write_content_size is not None:
1182 raise ValueError(
1246 raise ValueError(
@@ -1184,7 +1248,9 b' class ZstdCompressor(object):'
1184 )
1248 )
1185
1249
1186 if compression_params and write_dict_id is not None:
1250 if compression_params and write_dict_id is not None:
1187 raise ValueError("cannot define compression_params and " "write_dict_id")
1251 raise ValueError(
1252 "cannot define compression_params and " "write_dict_id"
1253 )
1188
1254
1189 if compression_params and threads:
1255 if compression_params and threads:
1190 raise ValueError("cannot define compression_params and threads")
1256 raise ValueError("cannot define compression_params and threads")
@@ -1201,7 +1267,9 b' class ZstdCompressor(object):'
1201
1267
1202 self._params = ffi.gc(params, lib.ZSTD_freeCCtxParams)
1268 self._params = ffi.gc(params, lib.ZSTD_freeCCtxParams)
1203
1269
1204 _set_compression_parameter(self._params, lib.ZSTD_c_compressionLevel, level)
1270 _set_compression_parameter(
1271 self._params, lib.ZSTD_c_compressionLevel, level
1272 )
1205
1273
1206 _set_compression_parameter(
1274 _set_compression_parameter(
1207 self._params,
1275 self._params,
@@ -1210,7 +1278,9 b' class ZstdCompressor(object):'
1210 )
1278 )
1211
1279
1212 _set_compression_parameter(
1280 _set_compression_parameter(
1213 self._params, lib.ZSTD_c_checksumFlag, 1 if write_checksum else 0
1281 self._params,
1282 lib.ZSTD_c_checksumFlag,
1283 1 if write_checksum else 0,
1214 )
1284 )
1215
1285
1216 _set_compression_parameter(
1286 _set_compression_parameter(
@@ -1218,7 +1288,9 b' class ZstdCompressor(object):'
1218 )
1288 )
1219
1289
1220 if threads:
1290 if threads:
1221 _set_compression_parameter(self._params, lib.ZSTD_c_nbWorkers, threads)
1291 _set_compression_parameter(
1292 self._params, lib.ZSTD_c_nbWorkers, threads
1293 )
1222
1294
1223 cctx = lib.ZSTD_createCCtx()
1295 cctx = lib.ZSTD_createCCtx()
1224 if cctx == ffi.NULL:
1296 if cctx == ffi.NULL:
@@ -1237,10 +1309,13 b' class ZstdCompressor(object):'
1237 )
1309 )
1238
1310
1239 def _setup_cctx(self):
1311 def _setup_cctx(self):
1240 zresult = lib.ZSTD_CCtx_setParametersUsingCCtxParams(self._cctx, self._params)
1312 zresult = lib.ZSTD_CCtx_setParametersUsingCCtxParams(
1313 self._cctx, self._params
1314 )
1241 if lib.ZSTD_isError(zresult):
1315 if lib.ZSTD_isError(zresult):
1242 raise ZstdError(
1316 raise ZstdError(
1243 "could not set compression parameters: %s" % _zstd_error(zresult)
1317 "could not set compression parameters: %s"
1318 % _zstd_error(zresult)
1244 )
1319 )
1245
1320
1246 dict_data = self._dict_data
1321 dict_data = self._dict_data
@@ -1259,7 +1334,8 b' class ZstdCompressor(object):'
1259
1334
1260 if lib.ZSTD_isError(zresult):
1335 if lib.ZSTD_isError(zresult):
1261 raise ZstdError(
1336 raise ZstdError(
1262 "could not load compression dictionary: %s" % _zstd_error(zresult)
1337 "could not load compression dictionary: %s"
1338 % _zstd_error(zresult)
1263 )
1339 )
1264
1340
1265 def memory_size(self):
1341 def memory_size(self):
@@ -1275,7 +1351,9 b' class ZstdCompressor(object):'
1275
1351
1276 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, len(data_buffer))
1352 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, len(data_buffer))
1277 if lib.ZSTD_isError(zresult):
1353 if lib.ZSTD_isError(zresult):
1278 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1354 raise ZstdError(
1355 "error setting source size: %s" % _zstd_error(zresult)
1356 )
1279
1357
1280 out_buffer = ffi.new("ZSTD_outBuffer *")
1358 out_buffer = ffi.new("ZSTD_outBuffer *")
1281 in_buffer = ffi.new("ZSTD_inBuffer *")
1359 in_buffer = ffi.new("ZSTD_inBuffer *")
@@ -1307,11 +1385,15 b' class ZstdCompressor(object):'
1307
1385
1308 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1386 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1309 if lib.ZSTD_isError(zresult):
1387 if lib.ZSTD_isError(zresult):
1310 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1388 raise ZstdError(
1389 "error setting source size: %s" % _zstd_error(zresult)
1390 )
1311
1391
1312 cobj = ZstdCompressionObj()
1392 cobj = ZstdCompressionObj()
1313 cobj._out = ffi.new("ZSTD_outBuffer *")
1393 cobj._out = ffi.new("ZSTD_outBuffer *")
1314 cobj._dst_buffer = ffi.new("char[]", COMPRESSION_RECOMMENDED_OUTPUT_SIZE)
1394 cobj._dst_buffer = ffi.new(
1395 "char[]", COMPRESSION_RECOMMENDED_OUTPUT_SIZE
1396 )
1315 cobj._out.dst = cobj._dst_buffer
1397 cobj._out.dst = cobj._dst_buffer
1316 cobj._out.size = COMPRESSION_RECOMMENDED_OUTPUT_SIZE
1398 cobj._out.size = COMPRESSION_RECOMMENDED_OUTPUT_SIZE
1317 cobj._out.pos = 0
1399 cobj._out.pos = 0
@@ -1328,7 +1410,9 b' class ZstdCompressor(object):'
1328
1410
1329 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1411 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1330 if lib.ZSTD_isError(zresult):
1412 if lib.ZSTD_isError(zresult):
1331 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1413 raise ZstdError(
1414 "error setting source size: %s" % _zstd_error(zresult)
1415 )
1332
1416
1333 return ZstdCompressionChunker(self, chunk_size=chunk_size)
1417 return ZstdCompressionChunker(self, chunk_size=chunk_size)
1334
1418
@@ -1353,7 +1437,9 b' class ZstdCompressor(object):'
1353
1437
1354 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1438 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1355 if lib.ZSTD_isError(zresult):
1439 if lib.ZSTD_isError(zresult):
1356 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1440 raise ZstdError(
1441 "error setting source size: %s" % _zstd_error(zresult)
1442 )
1357
1443
1358 in_buffer = ffi.new("ZSTD_inBuffer *")
1444 in_buffer = ffi.new("ZSTD_inBuffer *")
1359 out_buffer = ffi.new("ZSTD_outBuffer *")
1445 out_buffer = ffi.new("ZSTD_outBuffer *")
@@ -1381,7 +1467,9 b' class ZstdCompressor(object):'
1381 self._cctx, out_buffer, in_buffer, lib.ZSTD_e_continue
1467 self._cctx, out_buffer, in_buffer, lib.ZSTD_e_continue
1382 )
1468 )
1383 if lib.ZSTD_isError(zresult):
1469 if lib.ZSTD_isError(zresult):
1384 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
1470 raise ZstdError(
1471 "zstd compress error: %s" % _zstd_error(zresult)
1472 )
1385
1473
1386 if out_buffer.pos:
1474 if out_buffer.pos:
1387 ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
1475 ofh.write(ffi.buffer(out_buffer.dst, out_buffer.pos))
@@ -1423,7 +1511,9 b' class ZstdCompressor(object):'
1423
1511
1424 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1512 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1425 if lib.ZSTD_isError(zresult):
1513 if lib.ZSTD_isError(zresult):
1426 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1514 raise ZstdError(
1515 "error setting source size: %s" % _zstd_error(zresult)
1516 )
1427
1517
1428 return ZstdCompressionReader(self, source, read_size)
1518 return ZstdCompressionReader(self, source, read_size)
1429
1519
@@ -1443,7 +1533,9 b' class ZstdCompressor(object):'
1443 if size < 0:
1533 if size < 0:
1444 size = lib.ZSTD_CONTENTSIZE_UNKNOWN
1534 size = lib.ZSTD_CONTENTSIZE_UNKNOWN
1445
1535
1446 return ZstdCompressionWriter(self, writer, size, write_size, write_return_read)
1536 return ZstdCompressionWriter(
1537 self, writer, size, write_size, write_return_read
1538 )
1447
1539
1448 write_to = stream_writer
1540 write_to = stream_writer
1449
1541
@@ -1473,7 +1565,9 b' class ZstdCompressor(object):'
1473
1565
1474 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1566 zresult = lib.ZSTD_CCtx_setPledgedSrcSize(self._cctx, size)
1475 if lib.ZSTD_isError(zresult):
1567 if lib.ZSTD_isError(zresult):
1476 raise ZstdError("error setting source size: %s" % _zstd_error(zresult))
1568 raise ZstdError(
1569 "error setting source size: %s" % _zstd_error(zresult)
1570 )
1477
1571
1478 in_buffer = ffi.new("ZSTD_inBuffer *")
1572 in_buffer = ffi.new("ZSTD_inBuffer *")
1479 out_buffer = ffi.new("ZSTD_outBuffer *")
1573 out_buffer = ffi.new("ZSTD_outBuffer *")
@@ -1517,7 +1611,9 b' class ZstdCompressor(object):'
1517 self._cctx, out_buffer, in_buffer, lib.ZSTD_e_continue
1611 self._cctx, out_buffer, in_buffer, lib.ZSTD_e_continue
1518 )
1612 )
1519 if lib.ZSTD_isError(zresult):
1613 if lib.ZSTD_isError(zresult):
1520 raise ZstdError("zstd compress error: %s" % _zstd_error(zresult))
1614 raise ZstdError(
1615 "zstd compress error: %s" % _zstd_error(zresult)
1616 )
1521
1617
1522 if out_buffer.pos:
1618 if out_buffer.pos:
1523 data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
1619 data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
@@ -1596,10 +1692,14 b' def get_frame_parameters(data):'
1596 data_buffer = ffi.from_buffer(data)
1692 data_buffer = ffi.from_buffer(data)
1597 zresult = lib.ZSTD_getFrameHeader(params, data_buffer, len(data_buffer))
1693 zresult = lib.ZSTD_getFrameHeader(params, data_buffer, len(data_buffer))
1598 if lib.ZSTD_isError(zresult):
1694 if lib.ZSTD_isError(zresult):
1599 raise ZstdError("cannot get frame parameters: %s" % _zstd_error(zresult))
1695 raise ZstdError(
1696 "cannot get frame parameters: %s" % _zstd_error(zresult)
1697 )
1600
1698
1601 if zresult:
1699 if zresult:
1602 raise ZstdError("not enough data for frame parameters; need %d bytes" % zresult)
1700 raise ZstdError(
1701 "not enough data for frame parameters; need %d bytes" % zresult
1702 )
1603
1703
1604 return FrameParameters(params[0])
1704 return FrameParameters(params[0])
1605
1705
@@ -1611,9 +1711,14 b' class ZstdCompressionDict(object):'
1611 self.k = k
1711 self.k = k
1612 self.d = d
1712 self.d = d
1613
1713
1614 if dict_type not in (DICT_TYPE_AUTO, DICT_TYPE_RAWCONTENT, DICT_TYPE_FULLDICT):
1714 if dict_type not in (
1715 DICT_TYPE_AUTO,
1716 DICT_TYPE_RAWCONTENT,
1717 DICT_TYPE_FULLDICT,
1718 ):
1615 raise ValueError(
1719 raise ValueError(
1616 "invalid dictionary load mode: %d; must use " "DICT_TYPE_* constants"
1720 "invalid dictionary load mode: %d; must use "
1721 "DICT_TYPE_* constants"
1617 )
1722 )
1618
1723
1619 self._dict_type = dict_type
1724 self._dict_type = dict_type
@@ -1630,7 +1735,9 b' class ZstdCompressionDict(object):'
1630
1735
1631 def precompute_compress(self, level=0, compression_params=None):
1736 def precompute_compress(self, level=0, compression_params=None):
1632 if level and compression_params:
1737 if level and compression_params:
1633 raise ValueError("must only specify one of level or " "compression_params")
1738 raise ValueError(
1739 "must only specify one of level or " "compression_params"
1740 )
1634
1741
1635 if not level and not compression_params:
1742 if not level and not compression_params:
1636 raise ValueError("must specify one of level or compression_params")
1743 raise ValueError("must specify one of level or compression_params")
@@ -1675,7 +1782,9 b' class ZstdCompressionDict(object):'
1675 if ddict == ffi.NULL:
1782 if ddict == ffi.NULL:
1676 raise ZstdError("could not create decompression dict")
1783 raise ZstdError("could not create decompression dict")
1677
1784
1678 ddict = ffi.gc(ddict, lib.ZSTD_freeDDict, size=lib.ZSTD_sizeof_DDict(ddict))
1785 ddict = ffi.gc(
1786 ddict, lib.ZSTD_freeDDict, size=lib.ZSTD_sizeof_DDict(ddict)
1787 )
1679 self.__dict__["_ddict"] = ddict
1788 self.__dict__["_ddict"] = ddict
1680
1789
1681 return ddict
1790 return ddict
@@ -1805,7 +1914,9 b' class ZstdDecompressionObj(object):'
1805 self._decompressor._dctx, out_buffer, in_buffer
1914 self._decompressor._dctx, out_buffer, in_buffer
1806 )
1915 )
1807 if lib.ZSTD_isError(zresult):
1916 if lib.ZSTD_isError(zresult):
1808 raise ZstdError("zstd decompressor error: %s" % _zstd_error(zresult))
1917 raise ZstdError(
1918 "zstd decompressor error: %s" % _zstd_error(zresult)
1919 )
1809
1920
1810 if zresult == 0:
1921 if zresult == 0:
1811 self._finished = True
1922 self._finished = True
@@ -2105,16 +2216,22 b' class ZstdDecompressionReader(object):'
2105
2216
2106 if whence == os.SEEK_SET:
2217 if whence == os.SEEK_SET:
2107 if pos < 0:
2218 if pos < 0:
2108 raise ValueError("cannot seek to negative position with SEEK_SET")
2219 raise ValueError(
2220 "cannot seek to negative position with SEEK_SET"
2221 )
2109
2222
2110 if pos < self._bytes_decompressed:
2223 if pos < self._bytes_decompressed:
2111 raise ValueError("cannot seek zstd decompression stream " "backwards")
2224 raise ValueError(
2225 "cannot seek zstd decompression stream " "backwards"
2226 )
2112
2227
2113 read_amount = pos - self._bytes_decompressed
2228 read_amount = pos - self._bytes_decompressed
2114
2229
2115 elif whence == os.SEEK_CUR:
2230 elif whence == os.SEEK_CUR:
2116 if pos < 0:
2231 if pos < 0:
2117 raise ValueError("cannot seek zstd decompression stream " "backwards")
2232 raise ValueError(
2233 "cannot seek zstd decompression stream " "backwards"
2234 )
2118
2235
2119 read_amount = pos
2236 read_amount = pos
2120 elif whence == os.SEEK_END:
2237 elif whence == os.SEEK_END:
@@ -2123,7 +2240,9 b' class ZstdDecompressionReader(object):'
2123 )
2240 )
2124
2241
2125 while read_amount:
2242 while read_amount:
2126 result = self.read(min(read_amount, DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE))
2243 result = self.read(
2244 min(read_amount, DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE)
2245 )
2127
2246
2128 if not result:
2247 if not result:
2129 break
2248 break
@@ -2257,10 +2376,14 b' class ZstdDecompressionWriter(object):'
2257 while in_buffer.pos < in_buffer.size:
2376 while in_buffer.pos < in_buffer.size:
2258 zresult = lib.ZSTD_decompressStream(dctx, out_buffer, in_buffer)
2377 zresult = lib.ZSTD_decompressStream(dctx, out_buffer, in_buffer)
2259 if lib.ZSTD_isError(zresult):
2378 if lib.ZSTD_isError(zresult):
2260 raise ZstdError("zstd decompress error: %s" % _zstd_error(zresult))
2379 raise ZstdError(
2380 "zstd decompress error: %s" % _zstd_error(zresult)
2381 )
2261
2382
2262 if out_buffer.pos:
2383 if out_buffer.pos:
2263 self._writer.write(ffi.buffer(out_buffer.dst, out_buffer.pos)[:])
2384 self._writer.write(
2385 ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
2386 )
2264 total_write += out_buffer.pos
2387 total_write += out_buffer.pos
2265 out_buffer.pos = 0
2388 out_buffer.pos = 0
2266
2389
@@ -2299,7 +2422,9 b' class ZstdDecompressor(object):'
2299
2422
2300 data_buffer = ffi.from_buffer(data)
2423 data_buffer = ffi.from_buffer(data)
2301
2424
2302 output_size = lib.ZSTD_getFrameContentSize(data_buffer, len(data_buffer))
2425 output_size = lib.ZSTD_getFrameContentSize(
2426 data_buffer, len(data_buffer)
2427 )
2303
2428
2304 if output_size == lib.ZSTD_CONTENTSIZE_ERROR:
2429 if output_size == lib.ZSTD_CONTENTSIZE_ERROR:
2305 raise ZstdError("error determining content size from frame header")
2430 raise ZstdError("error determining content size from frame header")
@@ -2307,7 +2432,9 b' class ZstdDecompressor(object):'
2307 return b""
2432 return b""
2308 elif output_size == lib.ZSTD_CONTENTSIZE_UNKNOWN:
2433 elif output_size == lib.ZSTD_CONTENTSIZE_UNKNOWN:
2309 if not max_output_size:
2434 if not max_output_size:
2310 raise ZstdError("could not determine content size in frame header")
2435 raise ZstdError(
2436 "could not determine content size in frame header"
2437 )
2311
2438
2312 result_buffer = ffi.new("char[]", max_output_size)
2439 result_buffer = ffi.new("char[]", max_output_size)
2313 result_size = max_output_size
2440 result_size = max_output_size
@@ -2330,7 +2457,9 b' class ZstdDecompressor(object):'
2330 if lib.ZSTD_isError(zresult):
2457 if lib.ZSTD_isError(zresult):
2331 raise ZstdError("decompression error: %s" % _zstd_error(zresult))
2458 raise ZstdError("decompression error: %s" % _zstd_error(zresult))
2332 elif zresult:
2459 elif zresult:
2333 raise ZstdError("decompression error: did not decompress full frame")
2460 raise ZstdError(
2461 "decompression error: did not decompress full frame"
2462 )
2334 elif output_size and out_buffer.pos != output_size:
2463 elif output_size and out_buffer.pos != output_size:
2335 raise ZstdError(
2464 raise ZstdError(
2336 "decompression error: decompressed %d bytes; expected %d"
2465 "decompression error: decompressed %d bytes; expected %d"
@@ -2346,7 +2475,9 b' class ZstdDecompressor(object):'
2346 read_across_frames=False,
2475 read_across_frames=False,
2347 ):
2476 ):
2348 self._ensure_dctx()
2477 self._ensure_dctx()
2349 return ZstdDecompressionReader(self, source, read_size, read_across_frames)
2478 return ZstdDecompressionReader(
2479 self, source, read_size, read_across_frames
2480 )
2350
2481
2351 def decompressobj(self, write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE):
2482 def decompressobj(self, write_size=DECOMPRESSION_RECOMMENDED_OUTPUT_SIZE):
2352 if write_size < 1:
2483 if write_size < 1:
@@ -2421,9 +2552,13 b' class ZstdDecompressor(object):'
2421 while in_buffer.pos < in_buffer.size:
2552 while in_buffer.pos < in_buffer.size:
2422 assert out_buffer.pos == 0
2553 assert out_buffer.pos == 0
2423
2554
2424 zresult = lib.ZSTD_decompressStream(self._dctx, out_buffer, in_buffer)
2555 zresult = lib.ZSTD_decompressStream(
2556 self._dctx, out_buffer, in_buffer
2557 )
2425 if lib.ZSTD_isError(zresult):
2558 if lib.ZSTD_isError(zresult):
2426 raise ZstdError("zstd decompress error: %s" % _zstd_error(zresult))
2559 raise ZstdError(
2560 "zstd decompress error: %s" % _zstd_error(zresult)
2561 )
2427
2562
2428 if out_buffer.pos:
2563 if out_buffer.pos:
2429 data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
2564 data = ffi.buffer(out_buffer.dst, out_buffer.pos)[:]
@@ -2449,7 +2584,9 b' class ZstdDecompressor(object):'
2449 if not hasattr(writer, "write"):
2584 if not hasattr(writer, "write"):
2450 raise ValueError("must pass an object with a write() method")
2585 raise ValueError("must pass an object with a write() method")
2451
2586
2452 return ZstdDecompressionWriter(self, writer, write_size, write_return_read)
2587 return ZstdDecompressionWriter(
2588 self, writer, write_size, write_return_read
2589 )
2453
2590
2454 write_to = stream_writer
2591 write_to = stream_writer
2455
2592
@@ -2491,7 +2628,9 b' class ZstdDecompressor(object):'
2491
2628
2492 # Flush all read data to output.
2629 # Flush all read data to output.
2493 while in_buffer.pos < in_buffer.size:
2630 while in_buffer.pos < in_buffer.size:
2494 zresult = lib.ZSTD_decompressStream(self._dctx, out_buffer, in_buffer)
2631 zresult = lib.ZSTD_decompressStream(
2632 self._dctx, out_buffer, in_buffer
2633 )
2495 if lib.ZSTD_isError(zresult):
2634 if lib.ZSTD_isError(zresult):
2496 raise ZstdError(
2635 raise ZstdError(
2497 "zstd decompressor error: %s" % _zstd_error(zresult)
2636 "zstd decompressor error: %s" % _zstd_error(zresult)
@@ -2521,7 +2660,9 b' class ZstdDecompressor(object):'
2521 # All chunks should be zstd frames and should have content size set.
2660 # All chunks should be zstd frames and should have content size set.
2522 chunk_buffer = ffi.from_buffer(chunk)
2661 chunk_buffer = ffi.from_buffer(chunk)
2523 params = ffi.new("ZSTD_frameHeader *")
2662 params = ffi.new("ZSTD_frameHeader *")
2524 zresult = lib.ZSTD_getFrameHeader(params, chunk_buffer, len(chunk_buffer))
2663 zresult = lib.ZSTD_getFrameHeader(
2664 params, chunk_buffer, len(chunk_buffer)
2665 )
2525 if lib.ZSTD_isError(zresult):
2666 if lib.ZSTD_isError(zresult):
2526 raise ValueError("chunk 0 is not a valid zstd frame")
2667 raise ValueError("chunk 0 is not a valid zstd frame")
2527 elif zresult:
2668 elif zresult:
@@ -2546,7 +2687,9 b' class ZstdDecompressor(object):'
2546
2687
2547 zresult = lib.ZSTD_decompressStream(self._dctx, out_buffer, in_buffer)
2688 zresult = lib.ZSTD_decompressStream(self._dctx, out_buffer, in_buffer)
2548 if lib.ZSTD_isError(zresult):
2689 if lib.ZSTD_isError(zresult):
2549 raise ZstdError("could not decompress chunk 0: %s" % _zstd_error(zresult))
2690 raise ZstdError(
2691 "could not decompress chunk 0: %s" % _zstd_error(zresult)
2692 )
2550 elif zresult:
2693 elif zresult:
2551 raise ZstdError("chunk 0 did not decompress full frame")
2694 raise ZstdError("chunk 0 did not decompress full frame")
2552
2695
@@ -2561,11 +2704,15 b' class ZstdDecompressor(object):'
2561 raise ValueError("chunk %d must be bytes" % i)
2704 raise ValueError("chunk %d must be bytes" % i)
2562
2705
2563 chunk_buffer = ffi.from_buffer(chunk)
2706 chunk_buffer = ffi.from_buffer(chunk)
2564 zresult = lib.ZSTD_getFrameHeader(params, chunk_buffer, len(chunk_buffer))
2707 zresult = lib.ZSTD_getFrameHeader(
2708 params, chunk_buffer, len(chunk_buffer)
2709 )
2565 if lib.ZSTD_isError(zresult):
2710 if lib.ZSTD_isError(zresult):
2566 raise ValueError("chunk %d is not a valid zstd frame" % i)
2711 raise ValueError("chunk %d is not a valid zstd frame" % i)
2567 elif zresult:
2712 elif zresult:
2568 raise ValueError("chunk %d is too small to contain a zstd frame" % i)
2713 raise ValueError(
2714 "chunk %d is too small to contain a zstd frame" % i
2715 )
2569
2716
2570 if params.frameContentSize == lib.ZSTD_CONTENTSIZE_UNKNOWN:
2717 if params.frameContentSize == lib.ZSTD_CONTENTSIZE_UNKNOWN:
2571 raise ValueError("chunk %d missing content size in frame" % i)
2718 raise ValueError("chunk %d missing content size in frame" % i)
@@ -2580,7 +2727,9 b' class ZstdDecompressor(object):'
2580 in_buffer.size = len(chunk_buffer)
2727 in_buffer.size = len(chunk_buffer)
2581 in_buffer.pos = 0
2728 in_buffer.pos = 0
2582
2729
2583 zresult = lib.ZSTD_decompressStream(self._dctx, out_buffer, in_buffer)
2730 zresult = lib.ZSTD_decompressStream(
2731 self._dctx, out_buffer, in_buffer
2732 )
2584 if lib.ZSTD_isError(zresult):
2733 if lib.ZSTD_isError(zresult):
2585 raise ZstdError(
2734 raise ZstdError(
2586 "could not decompress chunk %d: %s" % _zstd_error(zresult)
2735 "could not decompress chunk %d: %s" % _zstd_error(zresult)
@@ -2597,7 +2746,9 b' class ZstdDecompressor(object):'
2597 lib.ZSTD_DCtx_reset(self._dctx, lib.ZSTD_reset_session_only)
2746 lib.ZSTD_DCtx_reset(self._dctx, lib.ZSTD_reset_session_only)
2598
2747
2599 if self._max_window_size:
2748 if self._max_window_size:
2600 zresult = lib.ZSTD_DCtx_setMaxWindowSize(self._dctx, self._max_window_size)
2749 zresult = lib.ZSTD_DCtx_setMaxWindowSize(
2750 self._dctx, self._max_window_size
2751 )
2601 if lib.ZSTD_isError(zresult):
2752 if lib.ZSTD_isError(zresult):
2602 raise ZstdError(
2753 raise ZstdError(
2603 "unable to set max window size: %s" % _zstd_error(zresult)
2754 "unable to set max window size: %s" % _zstd_error(zresult)
@@ -2605,11 +2756,14 b' class ZstdDecompressor(object):'
2605
2756
2606 zresult = lib.ZSTD_DCtx_setFormat(self._dctx, self._format)
2757 zresult = lib.ZSTD_DCtx_setFormat(self._dctx, self._format)
2607 if lib.ZSTD_isError(zresult):
2758 if lib.ZSTD_isError(zresult):
2608 raise ZstdError("unable to set decoding format: %s" % _zstd_error(zresult))
2759 raise ZstdError(
2760 "unable to set decoding format: %s" % _zstd_error(zresult)
2761 )
2609
2762
2610 if self._dict_data and load_dict:
2763 if self._dict_data and load_dict:
2611 zresult = lib.ZSTD_DCtx_refDDict(self._dctx, self._dict_data._ddict)
2764 zresult = lib.ZSTD_DCtx_refDDict(self._dctx, self._dict_data._ddict)
2612 if lib.ZSTD_isError(zresult):
2765 if lib.ZSTD_isError(zresult):
2613 raise ZstdError(
2766 raise ZstdError(
2614 "unable to reference prepared dictionary: %s" % _zstd_error(zresult)
2767 "unable to reference prepared dictionary: %s"
2768 % _zstd_error(zresult)
2615 )
2769 )
@@ -1,5 +1,5 b''
1 #require black
1 #require black
2
2
3 $ cd $RUNTESTDIR/..
3 $ cd $RUNTESTDIR/..
4 $ black --config=black.toml --check --diff `hg files 'set:(**.py + grep("^#!.*python")) - mercurial/thirdparty/** - "contrib/python-zstandard/**"'`
4 $ black --config=black.toml --check --diff `hg files 'set:(**.py + grep("^#!.*python")) - mercurial/thirdparty/**'`
5
5
General Comments 0
You need to be logged in to leave comments. Login now