contrib/python-zstandard/c-ext/compressoriterator.c
changeset 37495 b1fb341d8a61
parent 31796 e0dc40530c5a
child 42070 675775c33ab6
equal deleted inserted replaced
37494:1ce7a55b09d1 37495:b1fb341d8a61
    19 static void ZstdCompressorIterator_dealloc(ZstdCompressorIterator* self) {
    19 static void ZstdCompressorIterator_dealloc(ZstdCompressorIterator* self) {
    20 	Py_XDECREF(self->readResult);
    20 	Py_XDECREF(self->readResult);
    21 	Py_XDECREF(self->compressor);
    21 	Py_XDECREF(self->compressor);
    22 	Py_XDECREF(self->reader);
    22 	Py_XDECREF(self->reader);
    23 
    23 
    24 	if (self->buffer) {
    24 	if (self->buffer.buf) {
    25 		PyBuffer_Release(self->buffer);
    25 		PyBuffer_Release(&self->buffer);
    26 		PyMem_FREE(self->buffer);
    26 		memset(&self->buffer, 0, sizeof(self->buffer));
    27 		self->buffer = NULL;
       
    28 	}
    27 	}
    29 
    28 
    30 	if (self->output.dst) {
    29 	if (self->output.dst) {
    31 		PyMem_Free(self->output.dst);
    30 		PyMem_Free(self->output.dst);
    32 		self->output.dst = NULL;
    31 		self->output.dst = NULL;
    56 feedcompressor:
    55 feedcompressor:
    57 
    56 
    58 	/* If we have data left in the input, consume it. */
    57 	/* If we have data left in the input, consume it. */
    59 	if (self->input.pos < self->input.size) {
    58 	if (self->input.pos < self->input.size) {
    60 		Py_BEGIN_ALLOW_THREADS
    59 		Py_BEGIN_ALLOW_THREADS
    61 		if (self->compressor->mtcctx) {
    60 		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
    62 			zresult = ZSTDMT_compressStream(self->compressor->mtcctx,
    61 			&self->input, ZSTD_e_continue);
    63 				&self->output, &self->input);
       
    64 		}
       
    65 		else {
       
    66 			zresult = ZSTD_compressStream(self->compressor->cstream, &self->output,
       
    67 				&self->input);
       
    68 		}
       
    69 		Py_END_ALLOW_THREADS
    62 		Py_END_ALLOW_THREADS
    70 
    63 
    71 		/* Release the Python object holding the input buffer. */
    64 		/* Release the Python object holding the input buffer. */
    72 		if (self->input.pos == self->input.size) {
    65 		if (self->input.pos == self->input.size) {
    73 			self->input.src = NULL;
    66 			self->input.src = NULL;
   105 			}
    98 			}
   106 
    99 
   107 			PyBytes_AsStringAndSize(readResult, &readBuffer, &readSize);
   100 			PyBytes_AsStringAndSize(readResult, &readBuffer, &readSize);
   108 		}
   101 		}
   109 		else {
   102 		else {
   110 			assert(self->buffer && self->buffer->buf);
   103 			assert(self->buffer.buf);
   111 
   104 
   112 			/* Only support contiguous C arrays. */
   105 			/* Only support contiguous C arrays. */
   113 			assert(self->buffer->strides == NULL && self->buffer->suboffsets == NULL);
   106 			assert(self->buffer.strides == NULL && self->buffer.suboffsets == NULL);
   114 			assert(self->buffer->itemsize == 1);
   107 			assert(self->buffer.itemsize == 1);
   115 
   108 
   116 			readBuffer = (char*)self->buffer->buf + self->bufferOffset;
   109 			readBuffer = (char*)self->buffer.buf + self->bufferOffset;
   117 			bufferRemaining = self->buffer->len - self->bufferOffset;
   110 			bufferRemaining = self->buffer.len - self->bufferOffset;
   118 			readSize = min(bufferRemaining, (Py_ssize_t)self->inSize);
   111 			readSize = min(bufferRemaining, (Py_ssize_t)self->inSize);
   119 			self->bufferOffset += readSize;
   112 			self->bufferOffset += readSize;
   120 		}
   113 		}
   121 
   114 
   122 		if (0 == readSize) {
   115 		if (0 == readSize) {
   128 		}
   121 		}
   129 	}
   122 	}
   130 
   123 
   131 	/* EOF */
   124 	/* EOF */
   132 	if (0 == readSize) {
   125 	if (0 == readSize) {
   133 		if (self->compressor->mtcctx) {
   126 		self->input.src = NULL;
   134 			zresult = ZSTDMT_endStream(self->compressor->mtcctx, &self->output);
   127 		self->input.size = 0;
   135 		}
   128 		self->input.pos = 0;
   136 		else {
   129 
   137 			zresult = ZSTD_endStream(self->compressor->cstream, &self->output);
   130 		zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
   138 		}
   131 			&self->input, ZSTD_e_end);
   139 		if (ZSTD_isError(zresult)) {
   132 		if (ZSTD_isError(zresult)) {
   140 			PyErr_Format(ZstdError, "error ending compression stream: %s",
   133 			PyErr_Format(ZstdError, "error ending compression stream: %s",
   141 				ZSTD_getErrorName(zresult));
   134 				ZSTD_getErrorName(zresult));
   142 			return NULL;
   135 			return NULL;
   143 		}
   136 		}
   157 	self->input.src = readBuffer;
   150 	self->input.src = readBuffer;
   158 	self->input.size = readSize;
   151 	self->input.size = readSize;
   159 	self->input.pos = 0;
   152 	self->input.pos = 0;
   160 
   153 
   161 	Py_BEGIN_ALLOW_THREADS
   154 	Py_BEGIN_ALLOW_THREADS
   162 	if (self->compressor->mtcctx) {
   155 	zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
   163 		zresult = ZSTDMT_compressStream(self->compressor->mtcctx, &self->output,
   156 		&self->input, ZSTD_e_continue);
   164 			&self->input);
       
   165 	}
       
   166 	else {
       
   167 		zresult = ZSTD_compressStream(self->compressor->cstream, &self->output, &self->input);
       
   168 	}
       
   169 	Py_END_ALLOW_THREADS
   157 	Py_END_ALLOW_THREADS
   170 
   158 
   171 	/* The input buffer currently points to memory managed by Python
   159 	/* The input buffer currently points to memory managed by Python
   172 	(readBuffer). This object was allocated by this function. If it wasn't
   160 	(readBuffer). This object was allocated by this function. If it wasn't
   173 	fully consumed, we need to release it in a subsequent function call.
   161 	fully consumed, we need to release it in a subsequent function call.