contrib/python-zstandard/c-ext/compressionwriter.c
changeset 31796 e0dc40530c5a
parent 30895 c32454d69b85
child 37495 b1fb341d8a61
equal deleted inserted replaced
31795:2b130e26c3a4 31796:e0dc40530c5a
    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));