contrib/python-zstandard/c-ext/compressionreader.c

/**
* Copyright (c) 2017-present, Gregory Szorc
* All rights reserved.
*
* This software may be modified and distributed under the terms
* of the BSD license. See the LICENSE file for details.
*/

#include "python-zstandard.h"

extern PyObject* ZstdError;

static void set_unsupported_operation(void) {
	PyObject* iomod;
	PyObject* exc;

	iomod = PyImport_ImportModule("io");
	if (NULL == iomod) {
		return;
	}

	exc = PyObject_GetAttrString(iomod, "UnsupportedOperation");
	if (NULL == exc) {
		Py_DECREF(iomod);
		return;
	}

	PyErr_SetNone(exc);
	Py_DECREF(exc);
	Py_DECREF(iomod);
}

static void reader_dealloc(ZstdCompressionReader* self) {
	Py_XDECREF(self->compressor);
	Py_XDECREF(self->reader);

	if (self->buffer.buf) {
		PyBuffer_Release(&self->buffer);
		memset(&self->buffer, 0, sizeof(self->buffer));
	}

	PyObject_Del(self);
}

static ZstdCompressionReader* reader_enter(ZstdCompressionReader* self) {
	size_t zresult;

	if (self->entered) {
		PyErr_SetString(PyExc_ValueError, "cannot __enter__ multiple times");
		return NULL;
	}

	zresult = ZSTD_CCtx_setPledgedSrcSize(self->compressor->cctx, self->sourceSize);
	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "error setting source size: %s",
			ZSTD_getErrorName(zresult));
		return NULL;
	}

	self->entered = 1;

	Py_INCREF(self);
	return self;
}

static PyObject* reader_exit(ZstdCompressionReader* self, PyObject* args) {
	PyObject* exc_type;
	PyObject* exc_value;
	PyObject* exc_tb;

	if (!PyArg_ParseTuple(args, "OOO:__exit__", &exc_type, &exc_value, &exc_tb)) {
		return NULL;
	}

	self->entered = 0;
	self->closed = 1;

	/* Release resources associated with source. */
	Py_CLEAR(self->reader);
	if (self->buffer.buf) {
		PyBuffer_Release(&self->buffer);
		memset(&self->buffer, 0, sizeof(self->buffer));
	}

    Py_CLEAR(self->compressor);

	Py_RETURN_FALSE;
}

static PyObject* reader_readable(ZstdCompressionReader* self) {
	Py_RETURN_TRUE;
}

static PyObject* reader_writable(ZstdCompressionReader* self) {
	Py_RETURN_FALSE;
}

static PyObject* reader_seekable(ZstdCompressionReader* self) {
	Py_RETURN_FALSE;
}

static PyObject* reader_readline(PyObject* self, PyObject* args) {
	set_unsupported_operation();
	return NULL;
}

static PyObject* reader_readlines(PyObject* self, PyObject* args) {
	set_unsupported_operation();
	return NULL;
}

static PyObject* reader_write(PyObject* self, PyObject* args) {
	PyErr_SetString(PyExc_OSError, "stream is not writable");
	return NULL;
}

static PyObject* reader_writelines(PyObject* self, PyObject* args) {
	PyErr_SetString(PyExc_OSError, "stream is not writable");
	return NULL;
}

static PyObject* reader_isatty(PyObject* self) {
	Py_RETURN_FALSE;
}

static PyObject* reader_flush(PyObject* self) {
	Py_RETURN_NONE;
}

static PyObject* reader_close(ZstdCompressionReader* self) {
	self->closed = 1;
	Py_RETURN_NONE;
}

static PyObject* reader_closed(ZstdCompressionReader* self) {
	if (self->closed) {
		Py_RETURN_TRUE;
	}
	else {
		Py_RETURN_FALSE;
	}
}

static PyObject* reader_tell(ZstdCompressionReader* self) {
	/* TODO should this raise OSError since stream isn't seekable? */
	return PyLong_FromUnsignedLongLong(self->bytesCompressed);
}

static PyObject* reader_read(ZstdCompressionReader* self, PyObject* args, PyObject* kwargs) {
	static char* kwlist[] = {
		"size",
		NULL
	};

	Py_ssize_t size = -1;
	PyObject* result = NULL;
	char* resultBuffer;
	Py_ssize_t resultSize;
	size_t zresult;
	size_t oldPos;

	if (!self->entered) {
		PyErr_SetString(ZstdError, "read() must be called from an active context manager");
		return NULL;
	}

	if (self->closed) {
		PyErr_SetString(PyExc_ValueError, "stream is closed");
		return NULL;
	}

	if (self->finishedOutput) {
		return PyBytes_FromStringAndSize("", 0);
	}

	if (!PyArg_ParseTupleAndKeywords(args, kwargs, "n", kwlist, &size)) {
		return NULL;
	}

	if (size < 1) {
		PyErr_SetString(PyExc_ValueError, "cannot read negative or size 0 amounts");
		return NULL;
	}

	result = PyBytes_FromStringAndSize(NULL, size);
	if (NULL == result) {
		return NULL;
	}

	PyBytes_AsStringAndSize(result, &resultBuffer, &resultSize);

	self->output.dst = resultBuffer;
	self->output.size = resultSize;
	self->output.pos = 0;

readinput:

	/* If we have data left over, consume it. */
	if (self->input.pos < self->input.size) {
		oldPos = self->output.pos;

		Py_BEGIN_ALLOW_THREADS
		zresult = ZSTD_compress_generic(self->compressor->cctx,
			&self->output, &self->input, ZSTD_e_continue);

		Py_END_ALLOW_THREADS

		self->bytesCompressed += self->output.pos - oldPos;

		/* Input exhausted. Clear out state tracking. */
		if (self->input.pos == self->input.size) {
			memset(&self->input, 0, sizeof(self->input));
			Py_CLEAR(self->readResult);

			if (self->buffer.buf) {
				self->finishedInput = 1;
			}
		}

		if (ZSTD_isError(zresult)) {
			PyErr_Format(ZstdError, "zstd compress error: %s", ZSTD_getErrorName(zresult));
			return NULL;
		}

		if (self->output.pos) {
			/* If no more room in output, emit it. */
			if (self->output.pos == self->output.size) {
				memset(&self->output, 0, sizeof(self->output));
				return result;
			}

			/*
			 * There is room in the output. We fall through to below, which will either
			 * get more input for us or will attempt to end the stream.
			 */
		}

		/* Fall through to gather more input. */
	}

	if (!self->finishedInput) {
		if (self->reader) {
			Py_buffer buffer;

			assert(self->readResult == NULL);
			self->readResult = PyObject_CallMethod(self->reader, "read",
				"k", self->readSize);
			if (self->readResult == NULL) {
				return NULL;
			}

			memset(&buffer, 0, sizeof(buffer));

			if (0 != PyObject_GetBuffer(self->readResult, &buffer, PyBUF_CONTIG_RO)) {
				return NULL;
			}

			/* EOF */
			if (0 == buffer.len) {
				self->finishedInput = 1;
				Py_CLEAR(self->readResult);
			}
			else {
				self->input.src = buffer.buf;
				self->input.size = buffer.len;
				self->input.pos = 0;
			}

			PyBuffer_Release(&buffer);
		}
		else {
			assert(self->buffer.buf);

			self->input.src = self->buffer.buf;
			self->input.size = self->buffer.len;
			self->input.pos = 0;
		}
	}

	if (self->input.size) {
		goto readinput;
	}

	/* Else EOF */
	oldPos = self->output.pos;

	zresult = ZSTD_compress_generic(self->compressor->cctx, &self->output,
		&self->input, ZSTD_e_end);

	self->bytesCompressed += self->output.pos - oldPos;

	if (ZSTD_isError(zresult)) {
		PyErr_Format(ZstdError, "error ending compression stream: %s",
			ZSTD_getErrorName(zresult));
		return NULL;
	}

	assert(self->output.pos);

	if (0 == zresult) {
		self->finishedOutput = 1;
	}

	if (safe_pybytes_resize(&result, self->output.pos)) {
		Py_XDECREF(result);
		return NULL;
	}

	memset(&self->output, 0, sizeof(self->output));

	return result;
}

static PyObject* reader_readall(PyObject* self) {
	PyErr_SetNone(PyExc_NotImplementedError);
	return NULL;
}

static PyObject* reader_iter(PyObject* self) {
	set_unsupported_operation();
	return NULL;
}

static PyObject* reader_iternext(PyObject* self) {
	set_unsupported_operation();
	return NULL;
}

static PyMethodDef reader_methods[] = {
	{ "__enter__", (PyCFunction)reader_enter, METH_NOARGS,
	PyDoc_STR("Enter a compression context") },
	{ "__exit__", (PyCFunction)reader_exit, METH_VARARGS,
	PyDoc_STR("Exit a compression context") },
	{ "close", (PyCFunction)reader_close, METH_NOARGS,
	PyDoc_STR("Close the stream so it cannot perform any more operations") },
	{ "closed", (PyCFunction)reader_closed, METH_NOARGS,
	PyDoc_STR("Whether stream is closed") },
	{ "flush", (PyCFunction)reader_flush, METH_NOARGS, PyDoc_STR("no-ops") },
	{ "isatty", (PyCFunction)reader_isatty, METH_NOARGS, PyDoc_STR("Returns False") },
	{ "readable", (PyCFunction)reader_readable, METH_NOARGS,
	PyDoc_STR("Returns True") },
	{ "read", (PyCFunction)reader_read, METH_VARARGS | METH_KEYWORDS, PyDoc_STR("read compressed data") },
	{ "readall", (PyCFunction)reader_readall, METH_NOARGS, PyDoc_STR("Not implemented") },
	{ "readline", (PyCFunction)reader_readline, METH_VARARGS, PyDoc_STR("Not implemented") },
	{ "readlines", (PyCFunction)reader_readlines, METH_VARARGS, PyDoc_STR("Not implemented") },
	{ "seekable", (PyCFunction)reader_seekable, METH_NOARGS,
	PyDoc_STR("Returns False") },
	{ "tell", (PyCFunction)reader_tell, METH_NOARGS,
	PyDoc_STR("Returns current number of bytes compressed") },
	{ "writable", (PyCFunction)reader_writable, METH_NOARGS,
	PyDoc_STR("Returns False") },
	{ "write", reader_write, METH_VARARGS, PyDoc_STR("Raises OSError") },
	{ "writelines", reader_writelines, METH_VARARGS, PyDoc_STR("Not implemented") },
	{ NULL, NULL }
};

PyTypeObject ZstdCompressionReaderType = {
	PyVarObject_HEAD_INIT(NULL, 0)
	"zstd.ZstdCompressionReader", /* tp_name */
	sizeof(ZstdCompressionReader), /* tp_basicsize */
	0, /* tp_itemsize */
	(destructor)reader_dealloc, /* tp_dealloc */
	0, /* tp_print */
	0, /* tp_getattr */
	0, /* tp_setattr */
	0, /* tp_compare */
	0, /* tp_repr */
	0, /* tp_as_number */
	0, /* tp_as_sequence */
	0, /* tp_as_mapping */
	0, /* tp_hash */
	0, /* tp_call */
	0, /* tp_str */
	0, /* tp_getattro */
	0, /* tp_setattro */
	0, /* tp_as_buffer */
	Py_TPFLAGS_DEFAULT, /* tp_flags */
	0, /* tp_doc */
	0, /* tp_traverse */
	0, /* tp_clear */
	0, /* tp_richcompare */
	0, /* tp_weaklistoffset */
	reader_iter, /* tp_iter */
	reader_iternext, /* tp_iternext */
	reader_methods, /* tp_methods */
	0, /* tp_members */
	0, /* tp_getset */
	0, /* tp_base */
	0, /* tp_dict */
	0, /* tp_descr_get */
	0, /* tp_descr_set */
	0, /* tp_dictoffset */
	0, /* tp_init */
	0, /* tp_alloc */
	PyType_GenericNew, /* tp_new */
};

void compressionreader_module_init(PyObject* mod) {
	/* TODO make reader a sub-class of io.RawIOBase */

	Py_TYPE(&ZstdCompressionReaderType) = &PyType_Type;
	if (PyType_Ready(&ZstdCompressionReaderType) < 0) {
		return;
	}
}