113 int flushMode = compressorobj_flush_finish; |
113 int flushMode = compressorobj_flush_finish; |
114 size_t zresult; |
114 size_t zresult; |
115 PyObject* result = NULL; |
115 PyObject* result = NULL; |
116 Py_ssize_t resultSize = 0; |
116 Py_ssize_t resultSize = 0; |
117 ZSTD_inBuffer input; |
117 ZSTD_inBuffer input; |
|
118 ZSTD_EndDirective zFlushMode; |
118 |
119 |
119 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:flush", kwlist, &flushMode)) { |
120 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:flush", kwlist, &flushMode)) { |
120 return NULL; |
121 return NULL; |
121 } |
122 } |
122 |
123 |
126 } |
127 } |
127 |
128 |
128 if (self->finished) { |
129 if (self->finished) { |
129 PyErr_SetString(ZstdError, "compressor object already finished"); |
130 PyErr_SetString(ZstdError, "compressor object already finished"); |
130 return NULL; |
131 return NULL; |
|
132 } |
|
133 |
|
134 switch (flushMode) { |
|
135 case compressorobj_flush_block: |
|
136 zFlushMode = ZSTD_e_flush; |
|
137 break; |
|
138 |
|
139 case compressorobj_flush_finish: |
|
140 zFlushMode = ZSTD_e_end; |
|
141 self->finished = 1; |
|
142 break; |
|
143 |
|
144 default: |
|
145 PyErr_SetString(ZstdError, "unhandled flush mode"); |
|
146 return NULL; |
131 } |
147 } |
132 |
148 |
133 assert(self->output.pos == 0); |
149 assert(self->output.pos == 0); |
134 |
150 |
135 input.src = NULL; |
151 input.src = NULL; |
136 input.size = 0; |
152 input.size = 0; |
137 input.pos = 0; |
153 input.pos = 0; |
138 |
154 |
139 if (flushMode == compressorobj_flush_block) { |
155 while (1) { |
140 /* The output buffer is of size ZSTD_CStreamOutSize(), which is |
|
141 guaranteed to hold a full block. */ |
|
142 Py_BEGIN_ALLOW_THREADS |
156 Py_BEGIN_ALLOW_THREADS |
143 zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output, |
157 zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output, |
144 &input, ZSTD_e_flush); |
158 &input, zFlushMode); |
145 Py_END_ALLOW_THREADS |
159 Py_END_ALLOW_THREADS |
146 |
160 |
147 if (ZSTD_isError(zresult)) { |
|
148 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
|
149 return NULL; |
|
150 } |
|
151 |
|
152 /* Output buffer is guaranteed to hold full block. */ |
|
153 assert(zresult == 0); |
|
154 |
|
155 if (self->output.pos) { |
|
156 result = PyBytes_FromStringAndSize(self->output.dst, self->output.pos); |
|
157 if (!result) { |
|
158 return NULL; |
|
159 } |
|
160 } |
|
161 |
|
162 self->output.pos = 0; |
|
163 |
|
164 if (result) { |
|
165 return result; |
|
166 } |
|
167 else { |
|
168 return PyBytes_FromString(""); |
|
169 } |
|
170 } |
|
171 |
|
172 assert(flushMode == compressorobj_flush_finish); |
|
173 self->finished = 1; |
|
174 |
|
175 while (1) { |
|
176 zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output, |
|
177 &input, ZSTD_e_end); |
|
178 if (ZSTD_isError(zresult)) { |
161 if (ZSTD_isError(zresult)) { |
179 PyErr_Format(ZstdError, "error ending compression stream: %s", |
162 PyErr_Format(ZstdError, "error ending compression stream: %s", |
180 ZSTD_getErrorName(zresult)); |
163 ZSTD_getErrorName(zresult)); |
181 return NULL; |
164 return NULL; |
182 } |
165 } |