--- a/contrib/python-zstandard/c-ext/compressionwriter.c Sat Apr 01 13:43:52 2017 -0700
+++ b/contrib/python-zstandard/c-ext/compressionwriter.c Sat Apr 01 15:24:03 2017 -0700
@@ -18,11 +18,6 @@
Py_XDECREF(self->compressor);
Py_XDECREF(self->writer);
- if (self->cstream) {
- ZSTD_freeCStream(self->cstream);
- self->cstream = NULL;
- }
-
PyObject_Del(self);
}
@@ -32,9 +27,15 @@
return NULL;
}
- self->cstream = CStream_from_ZstdCompressor(self->compressor, self->sourceSize);
- if (!self->cstream) {
- return NULL;
+ if (self->compressor->mtcctx) {
+ if (init_mtcstream(self->compressor, self->sourceSize)) {
+ return NULL;
+ }
+ }
+ else {
+ if (0 != init_cstream(self->compressor, self->sourceSize)) {
+ return NULL;
+ }
}
self->entered = 1;
@@ -58,8 +59,8 @@
self->entered = 0;
- if (self->cstream && exc_type == Py_None && exc_value == Py_None &&
- exc_tb == Py_None) {
+ if ((self->compressor->cstream || self->compressor->mtcctx) && exc_type == Py_None
+ && exc_value == Py_None && exc_tb == Py_None) {
output.dst = PyMem_Malloc(self->outSize);
if (!output.dst) {
@@ -69,7 +70,12 @@
output.pos = 0;
while (1) {
- zresult = ZSTD_endStream(self->cstream, &output);
+ if (self->compressor->mtcctx) {
+ zresult = ZSTDMT_endStream(self->compressor->mtcctx, &output);
+ }
+ else {
+ zresult = ZSTD_endStream(self->compressor->cstream, &output);
+ }
if (ZSTD_isError(zresult)) {
PyErr_Format(ZstdError, "error ending compression stream: %s",
ZSTD_getErrorName(zresult));
@@ -95,21 +101,19 @@
}
PyMem_Free(output.dst);
- ZSTD_freeCStream(self->cstream);
- self->cstream = NULL;
}
Py_RETURN_FALSE;
}
static PyObject* ZstdCompressionWriter_memory_size(ZstdCompressionWriter* self) {
- if (!self->cstream) {
+ if (!self->compressor->cstream) {
PyErr_SetString(ZstdError, "cannot determine size of an inactive compressor; "
"call when a context manager is active");
return NULL;
}
- return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->cstream));
+ return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->compressor->cstream));
}
static PyObject* ZstdCompressionWriter_write(ZstdCompressionWriter* self, PyObject* args) {
@@ -147,7 +151,13 @@
while ((ssize_t)input.pos < sourceSize) {
Py_BEGIN_ALLOW_THREADS
- zresult = ZSTD_compressStream(self->cstream, &output, &input);
+ if (self->compressor->mtcctx) {
+ zresult = ZSTDMT_compressStream(self->compressor->mtcctx,
+ &output, &input);
+ }
+ else {
+ zresult = ZSTD_compressStream(self->compressor->cstream, &output, &input);
+ }
Py_END_ALLOW_THREADS
if (ZSTD_isError(zresult)) {
@@ -195,7 +205,12 @@
while (1) {
Py_BEGIN_ALLOW_THREADS
- zresult = ZSTD_flushStream(self->cstream, &output);
+ if (self->compressor->mtcctx) {
+ zresult = ZSTDMT_flushStream(self->compressor->mtcctx, &output);
+ }
+ else {
+ zresult = ZSTD_flushStream(self->compressor->cstream, &output);
+ }
Py_END_ALLOW_THREADS
if (ZSTD_isError(zresult)) {