equal
deleted
inserted
replaced
16 |
16 |
17 static void ZstdCompressionWriter_dealloc(ZstdCompressionWriter* self) { |
17 static void ZstdCompressionWriter_dealloc(ZstdCompressionWriter* self) { |
18 Py_XDECREF(self->compressor); |
18 Py_XDECREF(self->compressor); |
19 Py_XDECREF(self->writer); |
19 Py_XDECREF(self->writer); |
20 |
20 |
21 if (self->cstream) { |
|
22 ZSTD_freeCStream(self->cstream); |
|
23 self->cstream = NULL; |
|
24 } |
|
25 |
|
26 PyObject_Del(self); |
21 PyObject_Del(self); |
27 } |
22 } |
28 |
23 |
29 static PyObject* ZstdCompressionWriter_enter(ZstdCompressionWriter* self) { |
24 static PyObject* ZstdCompressionWriter_enter(ZstdCompressionWriter* self) { |
30 if (self->entered) { |
25 if (self->entered) { |
31 PyErr_SetString(ZstdError, "cannot __enter__ multiple times"); |
26 PyErr_SetString(ZstdError, "cannot __enter__ multiple times"); |
32 return NULL; |
27 return NULL; |
33 } |
28 } |
34 |
29 |
35 self->cstream = CStream_from_ZstdCompressor(self->compressor, self->sourceSize); |
30 if (self->compressor->mtcctx) { |
36 if (!self->cstream) { |
31 if (init_mtcstream(self->compressor, self->sourceSize)) { |
37 return NULL; |
32 return NULL; |
|
33 } |
|
34 } |
|
35 else { |
|
36 if (0 != init_cstream(self->compressor, self->sourceSize)) { |
|
37 return NULL; |
|
38 } |
38 } |
39 } |
39 |
40 |
40 self->entered = 1; |
41 self->entered = 1; |
41 |
42 |
42 Py_INCREF(self); |
43 Py_INCREF(self); |
56 return NULL; |
57 return NULL; |
57 } |
58 } |
58 |
59 |
59 self->entered = 0; |
60 self->entered = 0; |
60 |
61 |
61 if (self->cstream && exc_type == Py_None && exc_value == Py_None && |
62 if ((self->compressor->cstream || self->compressor->mtcctx) && exc_type == Py_None |
62 exc_tb == Py_None) { |
63 && exc_value == Py_None && exc_tb == Py_None) { |
63 |
64 |
64 output.dst = PyMem_Malloc(self->outSize); |
65 output.dst = PyMem_Malloc(self->outSize); |
65 if (!output.dst) { |
66 if (!output.dst) { |
66 return PyErr_NoMemory(); |
67 return PyErr_NoMemory(); |
67 } |
68 } |
68 output.size = self->outSize; |
69 output.size = self->outSize; |
69 output.pos = 0; |
70 output.pos = 0; |
70 |
71 |
71 while (1) { |
72 while (1) { |
72 zresult = ZSTD_endStream(self->cstream, &output); |
73 if (self->compressor->mtcctx) { |
|
74 zresult = ZSTDMT_endStream(self->compressor->mtcctx, &output); |
|
75 } |
|
76 else { |
|
77 zresult = ZSTD_endStream(self->compressor->cstream, &output); |
|
78 } |
73 if (ZSTD_isError(zresult)) { |
79 if (ZSTD_isError(zresult)) { |
74 PyErr_Format(ZstdError, "error ending compression stream: %s", |
80 PyErr_Format(ZstdError, "error ending compression stream: %s", |
75 ZSTD_getErrorName(zresult)); |
81 ZSTD_getErrorName(zresult)); |
76 PyMem_Free(output.dst); |
82 PyMem_Free(output.dst); |
77 return NULL; |
83 return NULL; |
93 |
99 |
94 output.pos = 0; |
100 output.pos = 0; |
95 } |
101 } |
96 |
102 |
97 PyMem_Free(output.dst); |
103 PyMem_Free(output.dst); |
98 ZSTD_freeCStream(self->cstream); |
|
99 self->cstream = NULL; |
|
100 } |
104 } |
101 |
105 |
102 Py_RETURN_FALSE; |
106 Py_RETURN_FALSE; |
103 } |
107 } |
104 |
108 |
105 static PyObject* ZstdCompressionWriter_memory_size(ZstdCompressionWriter* self) { |
109 static PyObject* ZstdCompressionWriter_memory_size(ZstdCompressionWriter* self) { |
106 if (!self->cstream) { |
110 if (!self->compressor->cstream) { |
107 PyErr_SetString(ZstdError, "cannot determine size of an inactive compressor; " |
111 PyErr_SetString(ZstdError, "cannot determine size of an inactive compressor; " |
108 "call when a context manager is active"); |
112 "call when a context manager is active"); |
109 return NULL; |
113 return NULL; |
110 } |
114 } |
111 |
115 |
112 return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->cstream)); |
116 return PyLong_FromSize_t(ZSTD_sizeof_CStream(self->compressor->cstream)); |
113 } |
117 } |
114 |
118 |
115 static PyObject* ZstdCompressionWriter_write(ZstdCompressionWriter* self, PyObject* args) { |
119 static PyObject* ZstdCompressionWriter_write(ZstdCompressionWriter* self, PyObject* args) { |
116 const char* source; |
120 const char* source; |
117 Py_ssize_t sourceSize; |
121 Py_ssize_t sourceSize; |
145 input.size = sourceSize; |
149 input.size = sourceSize; |
146 input.pos = 0; |
150 input.pos = 0; |
147 |
151 |
148 while ((ssize_t)input.pos < sourceSize) { |
152 while ((ssize_t)input.pos < sourceSize) { |
149 Py_BEGIN_ALLOW_THREADS |
153 Py_BEGIN_ALLOW_THREADS |
150 zresult = ZSTD_compressStream(self->cstream, &output, &input); |
154 if (self->compressor->mtcctx) { |
|
155 zresult = ZSTDMT_compressStream(self->compressor->mtcctx, |
|
156 &output, &input); |
|
157 } |
|
158 else { |
|
159 zresult = ZSTD_compressStream(self->compressor->cstream, &output, &input); |
|
160 } |
151 Py_END_ALLOW_THREADS |
161 Py_END_ALLOW_THREADS |
152 |
162 |
153 if (ZSTD_isError(zresult)) { |
163 if (ZSTD_isError(zresult)) { |
154 PyMem_Free(output.dst); |
164 PyMem_Free(output.dst); |
155 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
165 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
193 output.size = self->outSize; |
203 output.size = self->outSize; |
194 output.pos = 0; |
204 output.pos = 0; |
195 |
205 |
196 while (1) { |
206 while (1) { |
197 Py_BEGIN_ALLOW_THREADS |
207 Py_BEGIN_ALLOW_THREADS |
198 zresult = ZSTD_flushStream(self->cstream, &output); |
208 if (self->compressor->mtcctx) { |
|
209 zresult = ZSTDMT_flushStream(self->compressor->mtcctx, &output); |
|
210 } |
|
211 else { |
|
212 zresult = ZSTD_flushStream(self->compressor->cstream, &output); |
|
213 } |
199 Py_END_ALLOW_THREADS |
214 Py_END_ALLOW_THREADS |
200 |
215 |
201 if (ZSTD_isError(zresult)) { |
216 if (ZSTD_isError(zresult)) { |
202 PyMem_Free(output.dst); |
217 PyMem_Free(output.dst); |
203 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |
218 PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult)); |