|
1 /** |
|
2 * Copyright (c) 2016-present, Gregory Szorc |
|
3 * All rights reserved. |
|
4 * |
|
5 * This software may be modified and distributed under the terms |
|
6 * of the BSD license. See the LICENSE file for details. |
|
7 */ |
|
8 |
|
9 #include "python-zstandard.h" |
|
10 |
|
11 extern PyObject* ZstdError; |
|
12 |
|
13 PyDoc_STRVAR(ZstdCompresssionWriter__doc__, |
|
14 """A context manager used for writing compressed output to a writer.\n" |
|
15 ); |
|
16 |
|
17 static void ZstdCompressionWriter_dealloc(ZstdCompressionWriter* self) { |
|
18 Py_XDECREF(self->compressor); |
|
19 Py_XDECREF(self->writer); |
|
20 |
|
21 if (self->cstream) { |
|
22 ZSTD_freeCStream(self->cstream); |
|
23 self->cstream = NULL; |
|
24 } |
|
25 |
|
26 PyObject_Del(self); |
|
27 } |
|
28 |
|
29 static PyObject* ZstdCompressionWriter_enter(ZstdCompressionWriter* self) { |
|
30 if (self->entered) { |
|
31 PyErr_SetString(ZstdError, "cannot __enter__ multiple times"); |
|
32 return NULL; |
|
33 } |
|
34 |
|
35 self->cstream = CStream_from_ZstdCompressor(self->compressor, self->sourceSize); |
|
36 if (!self->cstream) { |
|
37 return NULL; |
|
38 } |
|
39 |
|
40 self->entered = 1; |
|
41 |
|
42 Py_INCREF(self); |
|
43 return (PyObject*)self; |
|
44 } |
|
45 |
|
46 static PyObject* ZstdCompressionWriter_exit(ZstdCompressionWriter* self, PyObject* args) { |
|
47 PyObject* exc_type; |
|
48 PyObject* exc_value; |
|
49 PyObject* exc_tb; |
|
50 size_t zresult; |
|
51 |
|
52 ZSTD_outBuffer output; |
|
53 PyObject* res; |
|
54 |
|
55 if (!PyArg_ParseTuple(args, "OOO", &exc_type, &exc_value, &exc_tb)) { |
|
56 return NULL; |
|
57 } |
|
58 |
|
59 self->entered = 0; |
|
60 |
|
61 if (self->cstream && exc_type == Py_None && exc_value == Py_None && |
|
62 exc_tb == Py_None) { |
|
63 |
|
64 output.dst = malloc(self->outSize); |
|
65 if (!output.dst) { |
|
66 return PyErr_NoMemory(); |
|
67 } |
|
68 output.size = self->outSize; |
|
69 output.pos = 0; |
|
70 |
|
71 while (1) { |
|
72 zresult = ZSTD_endStream(self->cstream, &output); |
|
73 if (ZSTD_isError(zresult)) { |
|
74 PyErr_Format(ZstdError, "error ending compression stream: %s", |
|
75 ZSTD_getErrorName(zresult)); |
|
76 free(output.dst); |
|
77 return NULL; |
|
78 } |
|
79 |
|
80 if (output.pos) { |
|
81 #if PY_MAJOR_VERSION >= 3 |
|
82 res = PyObject_CallMethod(self->writer, "write", "y#", |
|
83 #else |
|
84 res = PyObject_CallMethod(self->writer, "write", "s#", |
|
85 #endif |
|
86 output.dst, output.pos); |
|
87 Py_XDECREF(res); |
|
88 } |
|
89 |
|
90 if (!zresult) { |
|
91 break; |
|
92 } |
|
93 |
|
94 output.pos = 0; |
|
95 } |
|
96 |
|
97 free(output.dst); |
|
98 ZSTD_freeCStream(self->cstream); |
|
99 self->cstream = NULL; |
|
100 } |
|
101 |
|
102 Py_RETURN_FALSE; |
|
103 } |
|
104 |
|
105 static PyObject* ZstdCompressionWriter_memory_size(ZstdCompressionWriter* self) { |
|
106 if (!self->cstream) { |
|
107 PyErr_SetString(ZstdError, "cannot determine size of an inactive compressor; " |
|
108 "call when a context manager is active"); |
|
109 return NULL; |
|
110 } |
|
111 |
|
112 return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->cstream)); |
|
113 } |
|
114 |
|
115 static PyObject* ZstdCompressionWriter_write(ZstdCompressionWriter* self, PyObject* args) { |
|
116 const char* source; |
|
117 Py_ssize_t sourceSize; |
|
118 size_t zresult; |
|
119 ZSTD_inBuffer input; |
|
120 ZSTD_outBuffer output; |
|
121 PyObject* res; |
|
122 |
|
123 #if PY_MAJOR_VERSION >= 3 |
|
124 if (!PyArg_ParseTuple(args, "y#", &source, &sourceSize)) { |
|
125 #else |
|
126 if (!PyArg_ParseTuple(args, "s#", &source, &sourceSize)) { |
|
127 #endif |
|
128 return NULL; |
|
129 } |
|
130 |
|
131 if (!self->entered) { |
|
132 PyErr_SetString(ZstdError, "compress must be called from an active context manager"); |
|
133 return NULL; |
|
134 } |
|
135 |
|
136 output.dst = malloc(self->outSize); |
|
137 if (!output.dst) { |
|
138 return PyErr_NoMemory(); |
|
139 } |
|
140 output.size = self->outSize; |
|
141 output.pos = 0; |
|
142 |
|
143 input.src = source; |
|
144 input.size = sourceSize; |
|
145 input.pos = 0; |
|
146 |
|
147 while ((ssize_t)input.pos < sourceSize) { |
|
148 Py_BEGIN_ALLOW_THREADS |
|
149 zresult = ZSTD_compressStream(self->cstream, &output, &input); |
|
150 Py_END_ALLOW_THREADS |
|
151 |
|
152 if (ZSTD_isError(zresult)) { |
|
153 free(output.dst); |
|
154 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
|
155 return NULL; |
|
156 } |
|
157 |
|
158 /* Copy data from output buffer to writer. */ |
|
159 if (output.pos) { |
|
160 #if PY_MAJOR_VERSION >= 3 |
|
161 res = PyObject_CallMethod(self->writer, "write", "y#", |
|
162 #else |
|
163 res = PyObject_CallMethod(self->writer, "write", "s#", |
|
164 #endif |
|
165 output.dst, output.pos); |
|
166 Py_XDECREF(res); |
|
167 } |
|
168 output.pos = 0; |
|
169 } |
|
170 |
|
171 free(output.dst); |
|
172 |
|
173 /* TODO return bytes written */ |
|
174 Py_RETURN_NONE; |
|
175 } |
|
176 |
|
177 static PyMethodDef ZstdCompressionWriter_methods[] = { |
|
178 { "__enter__", (PyCFunction)ZstdCompressionWriter_enter, METH_NOARGS, |
|
179 PyDoc_STR("Enter a compression context.") }, |
|
180 { "__exit__", (PyCFunction)ZstdCompressionWriter_exit, METH_VARARGS, |
|
181 PyDoc_STR("Exit a compression context.") }, |
|
182 { "memory_size", (PyCFunction)ZstdCompressionWriter_memory_size, METH_NOARGS, |
|
183 PyDoc_STR("Obtain the memory size of the underlying compressor") }, |
|
184 { "write", (PyCFunction)ZstdCompressionWriter_write, METH_VARARGS, |
|
185 PyDoc_STR("Compress data") }, |
|
186 { NULL, NULL } |
|
187 }; |
|
188 |
|
189 PyTypeObject ZstdCompressionWriterType = { |
|
190 PyVarObject_HEAD_INIT(NULL, 0) |
|
191 "zstd.ZstdCompressionWriter", /* tp_name */ |
|
192 sizeof(ZstdCompressionWriter), /* tp_basicsize */ |
|
193 0, /* tp_itemsize */ |
|
194 (destructor)ZstdCompressionWriter_dealloc, /* tp_dealloc */ |
|
195 0, /* tp_print */ |
|
196 0, /* tp_getattr */ |
|
197 0, /* tp_setattr */ |
|
198 0, /* tp_compare */ |
|
199 0, /* tp_repr */ |
|
200 0, /* tp_as_number */ |
|
201 0, /* tp_as_sequence */ |
|
202 0, /* tp_as_mapping */ |
|
203 0, /* tp_hash */ |
|
204 0, /* tp_call */ |
|
205 0, /* tp_str */ |
|
206 0, /* tp_getattro */ |
|
207 0, /* tp_setattro */ |
|
208 0, /* tp_as_buffer */ |
|
209 Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */ |
|
210 ZstdCompresssionWriter__doc__, /* tp_doc */ |
|
211 0, /* tp_traverse */ |
|
212 0, /* tp_clear */ |
|
213 0, /* tp_richcompare */ |
|
214 0, /* tp_weaklistoffset */ |
|
215 0, /* tp_iter */ |
|
216 0, /* tp_iternext */ |
|
217 ZstdCompressionWriter_methods, /* tp_methods */ |
|
218 0, /* tp_members */ |
|
219 0, /* tp_getset */ |
|
220 0, /* tp_base */ |
|
221 0, /* tp_dict */ |
|
222 0, /* tp_descr_get */ |
|
223 0, /* tp_descr_set */ |
|
224 0, /* tp_dictoffset */ |
|
225 0, /* tp_init */ |
|
226 0, /* tp_alloc */ |
|
227 PyType_GenericNew, /* tp_new */ |
|
228 }; |
|
229 |
|
230 void compressionwriter_module_init(PyObject* mod) { |
|
231 Py_TYPE(&ZstdCompressionWriterType) = &PyType_Type; |
|
232 if (PyType_Ready(&ZstdCompressionWriterType) < 0) { |
|
233 return; |
|
234 } |
|
235 } |