34 ZSTD_inBuffer input; |
34 ZSTD_inBuffer input; |
35 size_t zresult; |
35 size_t zresult; |
36 PyObject* result = NULL; |
36 PyObject* result = NULL; |
37 Py_ssize_t resultSize = 0; |
37 Py_ssize_t resultSize = 0; |
38 |
38 |
39 if (self->flushed) { |
39 if (self->finished) { |
40 PyErr_SetString(ZstdError, "cannot call compress() after flush() has been called"); |
40 PyErr_SetString(ZstdError, "cannot call compress() after compressor finished"); |
41 return NULL; |
41 return NULL; |
42 } |
42 } |
43 |
43 |
44 #if PY_MAJOR_VERSION >= 3 |
44 #if PY_MAJOR_VERSION >= 3 |
45 if (!PyArg_ParseTuple(args, "y#", &source, &sourceSize)) { |
45 if (!PyArg_ParseTuple(args, "y#", &source, &sourceSize)) { |
90 else { |
90 else { |
91 return PyBytes_FromString(""); |
91 return PyBytes_FromString(""); |
92 } |
92 } |
93 } |
93 } |
94 |
94 |
95 static PyObject* ZstdCompressionObj_flush(ZstdCompressionObj* self) { |
95 static PyObject* ZstdCompressionObj_flush(ZstdCompressionObj* self, PyObject* args) { |
|
96 int flushMode = compressorobj_flush_finish; |
96 size_t zresult; |
97 size_t zresult; |
97 PyObject* result = NULL; |
98 PyObject* result = NULL; |
98 Py_ssize_t resultSize = 0; |
99 Py_ssize_t resultSize = 0; |
99 |
100 |
100 if (self->flushed) { |
101 if (!PyArg_ParseTuple(args, "|i", &flushMode)) { |
101 PyErr_SetString(ZstdError, "flush() already called"); |
102 return NULL; |
102 return NULL; |
103 } |
103 } |
104 |
104 |
105 if (flushMode != compressorobj_flush_finish && flushMode != compressorobj_flush_block) { |
105 self->flushed = 1; |
106 PyErr_SetString(PyExc_ValueError, "flush mode not recognized"); |
|
107 return NULL; |
|
108 } |
|
109 |
|
110 if (self->finished) { |
|
111 PyErr_SetString(ZstdError, "compressor object already finished"); |
|
112 return NULL; |
|
113 } |
|
114 |
|
115 assert(self->output.pos == 0); |
|
116 |
|
117 if (flushMode == compressorobj_flush_block) { |
|
118 /* The output buffer is of size ZSTD_CStreamOutSize(), which is |
|
119 guaranteed to hold a full block. */ |
|
120 Py_BEGIN_ALLOW_THREADS |
|
121 zresult = ZSTD_flushStream(self->cstream, &self->output); |
|
122 Py_END_ALLOW_THREADS |
|
123 |
|
124 if (ZSTD_isError(zresult)) { |
|
125 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
|
126 return NULL; |
|
127 } |
|
128 |
|
129 /* Output buffer is guaranteed to hold full block. */ |
|
130 assert(zresult == 0); |
|
131 |
|
132 if (self->output.pos) { |
|
133 result = PyBytes_FromStringAndSize(self->output.dst, self->output.pos); |
|
134 if (!result) { |
|
135 return NULL; |
|
136 } |
|
137 } |
|
138 |
|
139 self->output.pos = 0; |
|
140 |
|
141 if (result) { |
|
142 return result; |
|
143 } |
|
144 else { |
|
145 return PyBytes_FromString(""); |
|
146 } |
|
147 } |
|
148 |
|
149 assert(flushMode == compressorobj_flush_finish); |
|
150 self->finished = 1; |
106 |
151 |
107 while (1) { |
152 while (1) { |
108 zresult = ZSTD_endStream(self->cstream, &self->output); |
153 zresult = ZSTD_endStream(self->cstream, &self->output); |
109 if (ZSTD_isError(zresult)) { |
154 if (ZSTD_isError(zresult)) { |
110 PyErr_Format(ZstdError, "error ending compression stream: %s", |
155 PyErr_Format(ZstdError, "error ending compression stream: %s", |
149 } |
194 } |
150 |
195 |
151 static PyMethodDef ZstdCompressionObj_methods[] = { |
196 static PyMethodDef ZstdCompressionObj_methods[] = { |
152 { "compress", (PyCFunction)ZstdCompressionObj_compress, METH_VARARGS, |
197 { "compress", (PyCFunction)ZstdCompressionObj_compress, METH_VARARGS, |
153 PyDoc_STR("compress data") }, |
198 PyDoc_STR("compress data") }, |
154 { "flush", (PyCFunction)ZstdCompressionObj_flush, METH_NOARGS, |
199 { "flush", (PyCFunction)ZstdCompressionObj_flush, METH_VARARGS, |
155 PyDoc_STR("finish compression operation") }, |
200 PyDoc_STR("finish compression operation") }, |
156 { NULL, NULL } |
201 { NULL, NULL } |
157 }; |
202 }; |
158 |
203 |
159 PyTypeObject ZstdCompressionObjType = { |
204 PyTypeObject ZstdCompressionObjType = { |