diff --git a/c-ext/compressoriterator.c b/c-ext/compressoriterator.c index 98994b82..9f840b9a 100644 --- a/c-ext/compressoriterator.c +++ b/c-ext/compressoriterator.c @@ -62,7 +62,7 @@ static PyObject *ZstdCompressorIterator_iternext(ZstdCompressorIterator *self) { self->input.src = NULL; self->input.pos = 0; self->input.size = 0; - Py_DECREF(self->readResult); + Py_XDECREF(self->readResult); self->readResult = NULL; } diff --git a/tests/test_compressor_read_to_iter.py b/tests/test_compressor_read_to_iter.py index dd8a48f4..a8a59852 100644 --- a/tests/test_compressor_read_to_iter.py +++ b/tests/test_compressor_read_to_iter.py @@ -1,4 +1,5 @@ import io +import random import unittest import zstandard as zstd @@ -105,6 +106,40 @@ def test_read_write_size(self): self.assertEqual(source._read_count, len(source.getvalue()) + 1) + def test_read_buffer_small_output(self): + # Incompressible input spanning a block forces the iterator to resume + # with a partially consumed input buffer when output chunks are small. + data = random.Random(303).randbytes(131073) + expected = zstd.ZstdCompressor(level=3).compress(data) + + for source in ( + data, + bytearray(data), + memoryview(data), + io.BytesIO(data), + ): + for write_size in (1, 1024): + with self.subTest(source=type(source), write_size=write_size): + if hasattr(source, "seek"): + source.seek(0) + it = zstd.ZstdCompressor(level=3).read_to_iter( + source, + size=len(data), + read_size=len(data), + write_size=write_size, + ) + chunks = list(it) + self.assertGreater(len(chunks), 1) + self.assertTrue(all(len(c) <= write_size for c in chunks)) + compressed = b"".join(chunks) + self.assertEqual(compressed, expected) + self.assertEqual( + zstd.ZstdDecompressor().decompress(compressed), data + ) + for _ in range(2): + with self.assertRaises(StopIteration): + next(it) + def test_multithreaded(self): source = io.BytesIO() source.write(b"a" * 1048576)