contrib/python-zstandard/c-ext/compressobj.c
changeset 40121 73fef626dae3
parent 37495 b1fb341d8a61
child 42070 675775c33ab6
--- a/contrib/python-zstandard/c-ext/compressobj.c	Tue Sep 25 20:55:03 2018 +0900
+++ b/contrib/python-zstandard/c-ext/compressobj.c	Mon Oct 08 16:27:40 2018 -0700
@@ -115,6 +115,7 @@
 	PyObject* result = NULL;
 	Py_ssize_t resultSize = 0;
 	ZSTD_inBuffer input;
+	ZSTD_EndDirective zFlushMode;
 
 	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:flush", kwlist, &flushMode)) {
 		return NULL;
@@ -130,52 +131,34 @@
 		return NULL;
 	}
 
+	switch (flushMode) {
+		case compressorobj_flush_block:
+			zFlushMode = ZSTD_e_flush;
+			break;
+
+		case compressorobj_flush_finish:
+			zFlushMode = ZSTD_e_end;
+			self->finished = 1;
+			break;
+
+		default:
+			PyErr_SetString(ZstdError, "unhandled flush mode");
+			return NULL;
+	}
+
 	assert(self->output.pos == 0);
 
 	input.src = NULL;
 	input.size = 0;
 	input.pos = 0;
 
-	if (flushMode == compressorobj_flush_block) {
-		/* The output buffer is of size ZSTD_CStreamOutSize(), which is 
-		   guaranteed to hold a full block. */
+	while (1) {
 		Py_BEGIN_ALLOW_THREADS
-			zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
-				&input, ZSTD_e_flush);
+		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
+			&input, zFlushMode);
 		Py_END_ALLOW_THREADS
 
 		if (ZSTD_isError(zresult)) {
-			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
-			return NULL;
-		}
-
-		/* Output buffer is guaranteed to hold full block. */
-		assert(zresult == 0);
-
-		if (self->output.pos) {
-			result = PyBytes_FromStringAndSize(self->output.dst, self->output.pos);
-			if (!result) {
-				return NULL;
-			}
-		}
-
-		self->output.pos = 0;
-
-		if (result) {
-			return result;
-		}
-		else {
-			return PyBytes_FromString("");
-		}
-	}
-
-	assert(flushMode == compressorobj_flush_finish);
-	self->finished = 1;
-
-	while (1) {
-		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
-			&input, ZSTD_e_end);
-		if (ZSTD_isError(zresult)) {
 			PyErr_Format(ZstdError, "error ending compression stream: %s",
 				ZSTD_getErrorName(zresult));
 			return NULL;