71 while (1) { |
71 while (1) { |
72 zresult = ZSTD_endStream(self->cstream, &output); |
72 zresult = ZSTD_endStream(self->cstream, &output); |
73 if (ZSTD_isError(zresult)) { |
73 if (ZSTD_isError(zresult)) { |
74 PyErr_Format(ZstdError, "error ending compression stream: %s", |
74 PyErr_Format(ZstdError, "error ending compression stream: %s", |
75 ZSTD_getErrorName(zresult)); |
75 ZSTD_getErrorName(zresult)); |
76 free(output.dst); |
76 PyMem_Free(output.dst); |
77 return NULL; |
77 return NULL; |
78 } |
78 } |
79 |
79 |
80 if (output.pos) { |
80 if (output.pos) { |
81 #if PY_MAJOR_VERSION >= 3 |
81 #if PY_MAJOR_VERSION >= 3 |
148 Py_BEGIN_ALLOW_THREADS |
148 Py_BEGIN_ALLOW_THREADS |
149 zresult = ZSTD_compressStream(self->cstream, &output, &input); |
149 zresult = ZSTD_compressStream(self->cstream, &output, &input); |
150 Py_END_ALLOW_THREADS |
150 Py_END_ALLOW_THREADS |
151 |
151 |
152 if (ZSTD_isError(zresult)) { |
152 if (ZSTD_isError(zresult)) { |
153 free(output.dst); |
153 PyMem_Free(output.dst); |
154 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
154 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
155 return NULL; |
155 return NULL; |
156 } |
156 } |
157 |
157 |
158 /* Copy data from output buffer to writer. */ |
158 /* Copy data from output buffer to writer. */ |
166 Py_XDECREF(res); |
166 Py_XDECREF(res); |
167 } |
167 } |
168 output.pos = 0; |
168 output.pos = 0; |
169 } |
169 } |
170 |
170 |
171 free(output.dst); |
171 PyMem_Free(output.dst); |
172 |
172 |
173 /* TODO return bytes written */ |
173 /* TODO return bytes written */ |
174 Py_RETURN_NONE; |
174 Py_RETURN_NONE; |
175 } |
175 } |
|
176 |
|
177 static PyObject* ZstdCompressionWriter_flush(ZstdCompressionWriter* self, PyObject* args) { |
|
178 size_t zresult; |
|
179 ZSTD_outBuffer output; |
|
180 PyObject* res; |
|
181 |
|
182 if (!self->entered) { |
|
183 PyErr_SetString(ZstdError, "flush must be called from an active context manager"); |
|
184 return NULL; |
|
185 } |
|
186 |
|
187 output.dst = PyMem_Malloc(self->outSize); |
|
188 if (!output.dst) { |
|
189 return PyErr_NoMemory(); |
|
190 } |
|
191 output.size = self->outSize; |
|
192 output.pos = 0; |
|
193 |
|
194 while (1) { |
|
195 Py_BEGIN_ALLOW_THREADS |
|
196 zresult = ZSTD_flushStream(self->cstream, &output); |
|
197 Py_END_ALLOW_THREADS |
|
198 |
|
199 if (ZSTD_isError(zresult)) { |
|
200 PyMem_Free(output.dst); |
|
201 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
|
202 return NULL; |
|
203 } |
|
204 |
|
205 if (!output.pos) { |
|
206 break; |
|
207 } |
|
208 |
|
209 /* Copy data from output buffer to writer. */ |
|
210 if (output.pos) { |
|
211 #if PY_MAJOR_VERSION >= 3 |
|
212 res = PyObject_CallMethod(self->writer, "write", "y#", |
|
213 #else |
|
214 res = PyObject_CallMethod(self->writer, "write", "s#", |
|
215 #endif |
|
216 output.dst, output.pos); |
|
217 Py_XDECREF(res); |
|
218 } |
|
219 output.pos = 0; |
|
220 } |
|
221 |
|
222 PyMem_Free(output.dst); |
|
223 |
|
224 /* TODO return bytes written */ |
|
225 Py_RETURN_NONE; |
|
226 } |
176 |
227 |
177 static PyMethodDef ZstdCompressionWriter_methods[] = { |
228 static PyMethodDef ZstdCompressionWriter_methods[] = { |
178 { "__enter__", (PyCFunction)ZstdCompressionWriter_enter, METH_NOARGS, |
229 { "__enter__", (PyCFunction)ZstdCompressionWriter_enter, METH_NOARGS, |
179 PyDoc_STR("Enter a compression context.") }, |
230 PyDoc_STR("Enter a compression context.") }, |
180 { "__exit__", (PyCFunction)ZstdCompressionWriter_exit, METH_VARARGS, |
231 { "__exit__", (PyCFunction)ZstdCompressionWriter_exit, METH_VARARGS, |
181 PyDoc_STR("Exit a compression context.") }, |
232 PyDoc_STR("Exit a compression context.") }, |
182 { "memory_size", (PyCFunction)ZstdCompressionWriter_memory_size, METH_NOARGS, |
233 { "memory_size", (PyCFunction)ZstdCompressionWriter_memory_size, METH_NOARGS, |
183 PyDoc_STR("Obtain the memory size of the underlying compressor") }, |
234 PyDoc_STR("Obtain the memory size of the underlying compressor") }, |
184 { "write", (PyCFunction)ZstdCompressionWriter_write, METH_VARARGS, |
235 { "write", (PyCFunction)ZstdCompressionWriter_write, METH_VARARGS, |
185 PyDoc_STR("Compress data") }, |
236 PyDoc_STR("Compress data") }, |
|
237 { "flush", (PyCFunction)ZstdCompressionWriter_flush, METH_NOARGS, |
|
238 PyDoc_STR("Flush data and finish a zstd frame") }, |
186 { NULL, NULL } |
239 { NULL, NULL } |
187 }; |
240 }; |
188 |
241 |
189 PyTypeObject ZstdCompressionWriterType = { |
242 PyTypeObject ZstdCompressionWriterType = { |
190 PyVarObject_HEAD_INIT(NULL, 0) |
243 PyVarObject_HEAD_INIT(NULL, 0) |