<commit_before>/**
* Copyright (c) 2018-present, <NAME>
* All rights reserved.
*
* This software may be modified and distributed under the terms
* of the BSD license. See the LICENSE file for details.
*/

#include "python-zstandard.h"

extern PyObject* ZstdError;

PyDoc_STRVAR(ZstdCompressionChunkerIterator__doc__,
	"Iterator of output chunks from ZstdCompressionChunker.\n"
);

static void ZstdCompressionChunkerIterator_dealloc(ZstdCompressionChunkerIterator* self) {
	Py_XDECREF(self->chunker);

	PyObject_Del(self);
}

static PyObject* ZstdCompressionChunkerIterator_iter(PyObject* self) {
	Py_INCREF(self);
	return self;
}

static PyObject* ZstdCompressionChunkerIterator_iternext(ZstdCompressionChunkerIterator* self) {
	size_t zresult;
	PyObject* chunk;
	ZstdCompressionChunker* chunker = self->chunker;
	ZSTD_EndDirective zFlushMode;

	if (self->mode != compressionchunker_mode_normal && chunker->input.pos != chunker->input.size) {
		PyErr_SetString(ZstdError, "input should have been fully consumed before calling flush() or finish()");
		return NULL;
	}

	if (chunker->finished) {
		return NULL;
	}

	/* If we have data left in the input, consume it. */
	while (chunker->input.pos < chunker->input.size) {
		Py_BEGIN_ALLOW_THREADS
		zresult = ZSTD_compress_generic(chunker->compressor->cctx, &chunker->output,
			&chunker->input, ZSTD_e_continue);
		Py_END_ALLOW_THREADS

		/* Input is fully consumed. */
		if (chunker->input.pos == chunker->input.size) {
			chunker->input.src = NULL;
			chunker->input.pos = 0;
			chunker->input.size = 0;
			PyBuffer_Release(&chunker->inBuffer);
		}

		if (ZSTD_isError(zresult)) {
			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
			return NULL;
		}

		/* If it produced a full output chunk, emit it. */
		if (chunker->output.pos == chunker->output.size) {
			chunk = PyBytes_FromStringAndSize(chunker->output.dst, chunker->output.pos);
			if (!chunk) {
				return NULL;
			}

			chunker->output.pos = 0;

			return chunk;
		}

		/* Else continue to compress available input data. */
	}

	/* We also need this here for the special case of an empty input buffer. */
	if (chunker->input.pos == chunker->input.size) {
		chunker->input.src = NULL;
		chunker->input.pos = 0;
		chunker->input.size = 0;
		PyBuffer_Release(&chunker->inBuffer);
	}

	/* No more input data. A partial chunk may be in chunker->output.
	 * If we're in normal compression mode, we're done. Otherwise if we're in
	 * flush or finish mode, we need to emit what data remains.
	 */
	if (self->mode == compressionchunker_mode_normal) {
		/* We don't need to set StopIteration. */
		return NULL;
	}

	if (self->mode == compressionchunker_mode_flush) {
		zFlushMode = ZSTD_e_flush;
	}
	else if (self->mode == compressionchunker_mode_finish) {
		zFlushMode = ZSTD_e_end;
	}
	else {
		PyErr_SetString(ZstdError, "unhandled compression mode; this should never happen");
		return NULL;
	}

	Py_BEGIN_ALLOW_THREADS
	zresult = ZSTD_compress_generic(chunker->compressor->cctx, &chunker->output,
		&chunker->input, zFlushMode);
	Py_END_ALLOW_THREADS

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "zstd compress error: %s",
			ZSTD_getErrorName(zresult));
		return NULL;
	}

	if (!zresult && chunker->output.pos == 0) {
		return NULL;
	}

	chunk = PyBytes_FromStringAndSize(chunker->output.dst, chunker->output.pos);
	if (!chunk) {
		return NULL;
	}

	chunker->output.pos = 0;

	if (!zresult && self->mode == compressionchunker_mode_finish) {
		chunker->finished = 1;
	}

	return chunk;
}

PyTypeObject ZstdCompressionChunkerIteratorType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	"zstd.ZstdCompressionChunkerIterator", /* tp_name */
	sizeof(ZstdCompressionChunkerIterator), /* tp_basicsize */
	0,                               /* tp_itemsize */
	(destructor)ZstdCompressionChunkerIterator_dealloc, /* tp_dealloc */
	0,                               /* tp_print */
	0,                               /* tp_getattr */
	0,                               /* tp_setattr */
	0,                               /* tp_compare */
	0,                               /* tp_repr */
	0,                               /* tp_as_number */
	0,                               /* tp_as_sequence */
	0,                               /* tp_as_mapping */
	0,                               /* tp_hash */
	0,                               /* tp_call */
	0,                               /* tp_str */
	0,                               /* tp_getattro */
	0,                               /* tp_setattro */
	0,                               /* tp_as_buffer */
	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
	ZstdCompressionChunkerIterator__doc__, /* tp_doc */
	0,                               /* tp_traverse */
	0,                               /* tp_clear */
	0,                               /* tp_richcompare */
	0,                               /* tp_weaklistoffset */
	ZstdCompressionChunkerIterator_iter, /* tp_iter */
	(iternextfunc)ZstdCompressionChunkerIterator_iternext, /* tp_iternext */
	0,                               /* tp_methods */
	0,                               /* tp_members */
	0,                               /* tp_getset */
	0,                               /* tp_base */
	0,                               /* tp_dict */
	0,                               /* tp_descr_get */
	0,                               /* tp_descr_set */
	0,                               /* tp_dictoffset */
	0,                               /* tp_init */
	0,                               /* tp_alloc */
	PyType_GenericNew,              /* tp_new */
};

PyDoc_STRVAR(ZstdCompressionChunker__doc__,
	"Compress chunks iteratively into exact chunk sizes.\n"
);

static void ZstdCompressionChunker_dealloc(ZstdCompressionChunker* self) {
	PyMem_Free(self->output.dst);
	self->output.dst = NULL;
<commit_msg>compressionchunker: release Py_buffer when destroying chunker<commit_after>/**
* Copyright (c) 2018-present, <NAME>
* All rights reserved.
*
* This software may be modified and distributed under the terms
* of the BSD license. See the LICENSE file for details.
*/

#include "python-zstandard.h"

extern PyObject* ZstdError;

PyDoc_STRVAR(ZstdCompressionChunkerIterator__doc__,
	"Iterator of output chunks from ZstdCompressionChunker.\n"
);

static void ZstdCompressionChunkerIterator_dealloc(ZstdCompressionChunkerIterator* self) {
	Py_XDECREF(self->chunker);

	PyObject_Del(self);
}

static PyObject* ZstdCompressionChunkerIterator_iter(PyObject* self) {
	Py_INCREF(self);
	return self;
}

static PyObject* ZstdCompressionChunkerIterator_iternext(ZstdCompressionChunkerIterator* self) {
	size_t zresult;
	PyObject* chunk;
	ZstdCompressionChunker* chunker = self->chunker;
	ZSTD_EndDirective zFlushMode;

	if (self->mode != compressionchunker_mode_normal && chunker->input.pos != chunker->input.size) {
		PyErr_SetString(ZstdError, "input should have been fully consumed before calling flush() or finish()");
		return NULL;
	}

	if (chunker->finished) {
		return NULL;
	}

	/* If we have data left in the input, consume it. */
	while (chunker->input.pos < chunker->input.size) {
		Py_BEGIN_ALLOW_THREADS
		zresult = ZSTD_compress_generic(chunker->compressor->cctx, &chunker->output,
			&chunker->input, ZSTD_e_continue);
		Py_END_ALLOW_THREADS

		/* Input is fully consumed. */
		if (chunker->input.pos == chunker->input.size) {
			chunker->input.src = NULL;
			chunker->input.pos = 0;
			chunker->input.size = 0;
			PyBuffer_Release(&chunker->inBuffer);
		}

		if (ZSTD_isError(zresult)) {
			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
			return NULL;
		}

		/* If it produced a full output chunk, emit it. */
		if (chunker->output.pos == chunker->output.size) {
			chunk = PyBytes_FromStringAndSize(chunker->output.dst, chunker->output.pos);
			if (!chunk) {
				return NULL;
			}

			chunker->output.pos = 0;

			return chunk;
		}

		/* Else continue to compress available input data. */
	}

	/* We also need this here for the special case of an empty input buffer. */
	if (chunker->input.pos == chunker->input.size) {
		chunker->input.src = NULL;
		chunker->input.pos = 0;
		chunker->input.size = 0;
		PyBuffer_Release(&chunker->inBuffer);
	}

	/* No more input data. A partial chunk may be in chunker->output.
	 * If we're in normal compression mode, we're done. Otherwise if we're in
	 * flush or finish mode, we need to emit what data remains.
	 */
	if (self->mode == compressionchunker_mode_normal) {
		/* We don't need to set StopIteration. */
		return NULL;
	}

	if (self->mode == compressionchunker_mode_flush) {
		zFlushMode = ZSTD_e_flush;
	}
	else if (self->mode == compressionchunker_mode_finish) {
		zFlushMode = ZSTD_e_end;
	}
	else {
		PyErr_SetString(ZstdError, "unhandled compression mode; this should never happen");
		return NULL;
	}

	Py_BEGIN_ALLOW_THREADS
	zresult = ZSTD_compress_generic(chunker->compressor->cctx, &chunker->output,
		&chunker->input, zFlushMode);
	Py_END_ALLOW_THREADS

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "zstd compress error: %s",
			ZSTD_getErrorName(zresult));
		return NULL;
	}

	if (!zresult && chunker->output.pos == 0) {
		return NULL;
	}

	chunk = PyBytes_FromStringAndSize(chunker->output.dst, chunker->output.pos);
	if (!chunk) {
		return NULL;
	}

	chunker->output.pos = 0;

	if (!zresult && self->mode == compressionchunker_mode_finish) {
		chunker->finished = 1;
	}

	return chunk;
}

PyTypeObject ZstdCompressionChunkerIteratorType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	"zstd.ZstdCompressionChunkerIterator", /* tp_name */
	sizeof(ZstdCompressionChunkerIterator), /* tp_basicsize */
	0,                               /* tp_itemsize */
	(destructor)ZstdCompressionChunkerIterator_dealloc, /* tp_dealloc */
	0,                               /* tp_print */
	0,                               /* tp_getattr */
	0,                               /* tp_setattr */
	0,                               /* tp_compare */
	0,                               /* tp_repr */
	0,                               /* tp_as_number */
	0,                               /* tp_as_sequence */
	0,                               /* tp_as_mapping */
	0,                               /* tp_hash */
	0,                               /* tp_call */
	0,                               /* tp_str */
	0,                               /* tp_getattro */
	0,                               /* tp_setattro */
	0,                               /* tp_as_buffer */
	Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
	ZstdCompressionChunkerIterator__doc__, /* tp_doc */
	0,                               /* tp_traverse */
	0,                               /* tp_clear */
	0,                               /* tp_richcompare */
	0,                               /* tp_weaklistoffset */
	ZstdCompressionChunkerIterator_iter, /* tp_iter */
	(iternextfunc)ZstdCompressionChunkerIterator_iternext, /* tp_iternext */
	0,                               /* tp_methods */
	0,                               /* tp_members */
	0,                               /* tp_getset */
	0,                               /* tp_base */
	0,                               /* tp_dict */
	0,                               /* tp_descr_get */
	0,                               /* tp_descr_set */
	0,                               /* tp_dictoffset */
	0,                               /* tp_init */
	0,                               /* tp_alloc */
	PyType_GenericNew,              /* tp_new */
};

PyDoc_STRVAR(ZstdCompressionChunker__doc__,
	"Compress chunks iteratively into exact chunk sizes.\n"
);

static void ZstdCompressionChunker_dealloc(ZstdCompressionChunker* self) {
	PyBuffer_Release(&self->inBuffer);
	self->input.src = NULL;

	PyMem_Free(self->output.dst);
	self->output.dst = NULL;

{'fraction_non_alphanumeric': 0.11288677320639239, 'fraction_numerical': 0.007140428425705543, 'mean_word_length': 2.1711590296495955, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 97, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 58, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7289007', 'n_tokens_mistral': 3920, 'n_tokens_neox': 3774, 'n_words': 784}