--- a/contrib/python-zstandard/c-ext/compressionwriter.c Thu Apr 04 15:24:03 2019 -0700
+++ b/contrib/python-zstandard/c-ext/compressionwriter.c Thu Apr 04 17:34:43 2019 -0700
@@ -18,24 +18,23 @@
Py_XDECREF(self->compressor);
Py_XDECREF(self->writer);
+ PyMem_Free(self->output.dst);
+ self->output.dst = NULL;
+
PyObject_Del(self);
}
static PyObject* ZstdCompressionWriter_enter(ZstdCompressionWriter* self) {
- size_t zresult;
+ if (self->closed) {
+ PyErr_SetString(PyExc_ValueError, "stream is closed");
+ return NULL;
+ }
if (self->entered) {
PyErr_SetString(ZstdError, "cannot __enter__ multiple times");
return NULL;
}
- zresult = ZSTD_CCtx_setPledgedSrcSize(self->compressor->cctx, self->sourceSize);
- if (ZSTD_isError(zresult)) {
- PyErr_Format(ZstdError, "error setting source size: %s",
- ZSTD_getErrorName(zresult));
- return NULL;
- }
-
self->entered = 1;
Py_INCREF(self);
@@ -46,10 +45,6 @@
PyObject* exc_type;
PyObject* exc_value;
PyObject* exc_tb;
- size_t zresult;
-
- ZSTD_outBuffer output;
- PyObject* res;
if (!PyArg_ParseTuple(args, "OOO:__exit__", &exc_type, &exc_value, &exc_tb)) {
return NULL;
@@ -58,46 +53,11 @@
self->entered = 0;
if (exc_type == Py_None && exc_value == Py_None && exc_tb == Py_None) {
- ZSTD_inBuffer inBuffer;
-
- inBuffer.src = NULL;
- inBuffer.size = 0;
- inBuffer.pos = 0;
-
- output.dst = PyMem_Malloc(self->outSize);
- if (!output.dst) {
- return PyErr_NoMemory();
- }
- output.size = self->outSize;
- output.pos = 0;
+ PyObject* result = PyObject_CallMethod((PyObject*)self, "close", NULL);
- while (1) {
- zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &inBuffer, ZSTD_e_end);
- if (ZSTD_isError(zresult)) {
- PyErr_Format(ZstdError, "error ending compression stream: %s",
- ZSTD_getErrorName(zresult));
- PyMem_Free(output.dst);
- return NULL;
- }
-
- if (output.pos) {
-#if PY_MAJOR_VERSION >= 3
- res = PyObject_CallMethod(self->writer, "write", "y#",
-#else
- res = PyObject_CallMethod(self->writer, "write", "s#",
-#endif
- output.dst, output.pos);
- Py_XDECREF(res);
- }
-
- if (!zresult) {
- break;
- }
-
- output.pos = 0;
+ if (NULL == result) {
+ return NULL;
}
-
- PyMem_Free(output.dst);
}
Py_RETURN_FALSE;
@@ -117,7 +77,6 @@
Py_buffer source;
size_t zresult;
ZSTD_inBuffer input;
- ZSTD_outBuffer output;
PyObject* res;
Py_ssize_t totalWrite = 0;
@@ -130,143 +89,240 @@
return NULL;
}
- if (!self->entered) {
- PyErr_SetString(ZstdError, "compress must be called from an active context manager");
- goto finally;
- }
-
if (!PyBuffer_IsContiguous(&source, 'C') || source.ndim > 1) {
PyErr_SetString(PyExc_ValueError,
"data buffer should be contiguous and have at most one dimension");
goto finally;
}
- output.dst = PyMem_Malloc(self->outSize);
- if (!output.dst) {
- PyErr_NoMemory();
- goto finally;
+ if (self->closed) {
+ PyErr_SetString(PyExc_ValueError, "stream is closed");
+ return NULL;
}
- output.size = self->outSize;
- output.pos = 0;
+
+ self->output.pos = 0;
input.src = source.buf;
input.size = source.len;
input.pos = 0;
- while ((ssize_t)input.pos < source.len) {
+ while (input.pos < (size_t)source.len) {
Py_BEGIN_ALLOW_THREADS
- zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &input, ZSTD_e_continue);
+ zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output, &input, ZSTD_e_continue);
Py_END_ALLOW_THREADS
if (ZSTD_isError(zresult)) {
- PyMem_Free(output.dst);
PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
goto finally;
}
/* Copy data from output buffer to writer. */
- if (output.pos) {
+ if (self->output.pos) {
#if PY_MAJOR_VERSION >= 3
res = PyObject_CallMethod(self->writer, "write", "y#",
#else
res = PyObject_CallMethod(self->writer, "write", "s#",
#endif
- output.dst, output.pos);
+ self->output.dst, self->output.pos);
Py_XDECREF(res);
- totalWrite += output.pos;
- self->bytesCompressed += output.pos;
+ totalWrite += self->output.pos;
+ self->bytesCompressed += self->output.pos;
}
- output.pos = 0;
+ self->output.pos = 0;
}
- PyMem_Free(output.dst);
-
- result = PyLong_FromSsize_t(totalWrite);
+ if (self->writeReturnRead) {
+ result = PyLong_FromSize_t(input.pos);
+ }
+ else {
+ result = PyLong_FromSsize_t(totalWrite);
+ }
finally:
PyBuffer_Release(&source);
return result;
}
-static PyObject* ZstdCompressionWriter_flush(ZstdCompressionWriter* self, PyObject* args) {
+static PyObject* ZstdCompressionWriter_flush(ZstdCompressionWriter* self, PyObject* args, PyObject* kwargs) {
+ static char* kwlist[] = {
+ "flush_mode",
+ NULL
+ };
+
size_t zresult;
- ZSTD_outBuffer output;
ZSTD_inBuffer input;
PyObject* res;
Py_ssize_t totalWrite = 0;
+ unsigned flush_mode = 0;
+ ZSTD_EndDirective flush;
- if (!self->entered) {
- PyErr_SetString(ZstdError, "flush must be called from an active context manager");
+ if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|I:flush",
+ kwlist, &flush_mode)) {
return NULL;
}
+ switch (flush_mode) {
+ case 0:
+ flush = ZSTD_e_flush;
+ break;
+ case 1:
+ flush = ZSTD_e_end;
+ break;
+ default:
+ PyErr_Format(PyExc_ValueError, "unknown flush_mode: %d", flush_mode);
+ return NULL;
+ }
+
+ if (self->closed) {
+ PyErr_SetString(PyExc_ValueError, "stream is closed");
+ return NULL;
+ }
+
+ self->output.pos = 0;
+
input.src = NULL;
input.size = 0;
input.pos = 0;
- output.dst = PyMem_Malloc(self->outSize);
- if (!output.dst) {
- return PyErr_NoMemory();
- }
- output.size = self->outSize;
- output.pos = 0;
-
while (1) {
Py_BEGIN_ALLOW_THREADS
- zresult = ZSTD_compress_generic(self->compressor->cctx, &output, &input, ZSTD_e_flush);
+ zresult = ZSTD_compressStream2(self->compressor->cctx, &self->output, &input, flush);
Py_END_ALLOW_THREADS
if (ZSTD_isError(zresult)) {
- PyMem_Free(output.dst);
PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
return NULL;
}
/* Copy data from output buffer to writer. */
- if (output.pos) {
+ if (self->output.pos) {
#if PY_MAJOR_VERSION >= 3
res = PyObject_CallMethod(self->writer, "write", "y#",
#else
res = PyObject_CallMethod(self->writer, "write", "s#",
#endif
- output.dst, output.pos);
+ self->output.dst, self->output.pos);
Py_XDECREF(res);
- totalWrite += output.pos;
- self->bytesCompressed += output.pos;
+ totalWrite += self->output.pos;
+ self->bytesCompressed += self->output.pos;
}
- output.pos = 0;
+ self->output.pos = 0;
if (!zresult) {
break;
}
}
- PyMem_Free(output.dst);
+ return PyLong_FromSsize_t(totalWrite);
+}
+
+static PyObject* ZstdCompressionWriter_close(ZstdCompressionWriter* self) {
+ PyObject* result;
+
+ if (self->closed) {
+ Py_RETURN_NONE;
+ }
+
+ result = PyObject_CallMethod((PyObject*)self, "flush", "I", 1);
+ self->closed = 1;
+
+ if (NULL == result) {
+ return NULL;
+ }
- return PyLong_FromSsize_t(totalWrite);
+ /* Call close on underlying stream as well. */
+ if (PyObject_HasAttrString(self->writer, "close")) {
+ return PyObject_CallMethod(self->writer, "close", NULL);
+ }
+
+ Py_RETURN_NONE;
+}
+
+static PyObject* ZstdCompressionWriter_fileno(ZstdCompressionWriter* self) {
+ if (PyObject_HasAttrString(self->writer, "fileno")) {
+ return PyObject_CallMethod(self->writer, "fileno", NULL);
+ }
+ else {
+ PyErr_SetString(PyExc_OSError, "fileno not available on underlying writer");
+ return NULL;
+ }
}
static PyObject* ZstdCompressionWriter_tell(ZstdCompressionWriter* self) {
return PyLong_FromUnsignedLongLong(self->bytesCompressed);
}
+static PyObject* ZstdCompressionWriter_writelines(PyObject* self, PyObject* args) {
+ PyErr_SetNone(PyExc_NotImplementedError);
+ return NULL;
+}
+
+static PyObject* ZstdCompressionWriter_false(PyObject* self, PyObject* args) {
+ Py_RETURN_FALSE;
+}
+
+static PyObject* ZstdCompressionWriter_true(PyObject* self, PyObject* args) {
+ Py_RETURN_TRUE;
+}
+
+static PyObject* ZstdCompressionWriter_unsupported(PyObject* self, PyObject* args, PyObject* kwargs) {
+ PyObject* iomod;
+ PyObject* exc;
+
+ iomod = PyImport_ImportModule("io");
+ if (NULL == iomod) {
+ return NULL;
+ }
+
+ exc = PyObject_GetAttrString(iomod, "UnsupportedOperation");
+ if (NULL == exc) {
+ Py_DECREF(iomod);
+ return NULL;
+ }
+
+ PyErr_SetNone(exc);
+ Py_DECREF(exc);
+ Py_DECREF(iomod);
+
+ return NULL;
+}
+
static PyMethodDef ZstdCompressionWriter_methods[] = {
{ "__enter__", (PyCFunction)ZstdCompressionWriter_enter, METH_NOARGS,
PyDoc_STR("Enter a compression context.") },
{ "__exit__", (PyCFunction)ZstdCompressionWriter_exit, METH_VARARGS,
PyDoc_STR("Exit a compression context.") },
+ { "close", (PyCFunction)ZstdCompressionWriter_close, METH_NOARGS, NULL },
+ { "fileno", (PyCFunction)ZstdCompressionWriter_fileno, METH_NOARGS, NULL },
+ { "isatty", (PyCFunction)ZstdCompressionWriter_false, METH_NOARGS, NULL },
+ { "readable", (PyCFunction)ZstdCompressionWriter_false, METH_NOARGS, NULL },
+ { "readline", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "readlines", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "seek", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "seekable", ZstdCompressionWriter_false, METH_NOARGS, NULL },
+ { "truncate", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "writable", ZstdCompressionWriter_true, METH_NOARGS, NULL },
+ { "writelines", ZstdCompressionWriter_writelines, METH_VARARGS, NULL },
+ { "read", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "readall", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
+ { "readinto", (PyCFunction)ZstdCompressionWriter_unsupported, METH_VARARGS | METH_KEYWORDS, NULL },
{ "memory_size", (PyCFunction)ZstdCompressionWriter_memory_size, METH_NOARGS,
PyDoc_STR("Obtain the memory size of the underlying compressor") },
{ "write", (PyCFunction)ZstdCompressionWriter_write, METH_VARARGS | METH_KEYWORDS,
PyDoc_STR("Compress data") },
- { "flush", (PyCFunction)ZstdCompressionWriter_flush, METH_NOARGS,
+ { "flush", (PyCFunction)ZstdCompressionWriter_flush, METH_VARARGS | METH_KEYWORDS,
PyDoc_STR("Flush data and finish a zstd frame") },
{ "tell", (PyCFunction)ZstdCompressionWriter_tell, METH_NOARGS,
PyDoc_STR("Returns current number of bytes compressed") },
{ NULL, NULL }
};
+static PyMemberDef ZstdCompressionWriter_members[] = {
+ { "closed", T_BOOL, offsetof(ZstdCompressionWriter, closed), READONLY, NULL },
+ { NULL }
+};
+
PyTypeObject ZstdCompressionWriterType = {
PyVarObject_HEAD_INIT(NULL, 0)
"zstd.ZstdCompressionWriter", /* tp_name */
@@ -296,7 +352,7 @@
0, /* tp_iter */
0, /* tp_iternext */
ZstdCompressionWriter_methods, /* tp_methods */
- 0, /* tp_members */
+ ZstdCompressionWriter_members, /* tp_members */
0, /* tp_getset */
0, /* tp_base */
0, /* tp_dict */