Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
150 lines (120 sloc) 4.94 KB
#define BLOCK_SIZE (1024 * 128)
#include <python2.6/Python.h>
#include <lzo/lzo1x.h>
#include <stdio.h>
/* The following formula gives the worst possible compressed size. */
#define lzo1x_worst_compress(x) ((x) + ((x) / 16) + 64 + 3)
/* The size of and pointer to the shared buffer. */
static int block_size, buffer_size = 0;
static unsigned char *shared_buffer = NULL;
/* Don't store the size of compressed blocks in headers and trust the user to
* configure the correct block size? */
static int omit_headers = 0;
/* Working memory required by LZO library, allocated on first use. */
static char *working_memory = NULL;
#define ADD_SIZE(p) (omit_headers ? (p) : ((p) + sizeof(int)))
#define SUB_SIZE(p) (omit_headers ? (p) : ((p) - sizeof(int)))
static unsigned char *
get_buffer(int length)
{
if (omit_headers) {
if (!shared_buffer)
shared_buffer = malloc(buffer_size);
} else if (!shared_buffer || buffer_size < length) {
free(shared_buffer);
shared_buffer = malloc(length);
buffer_size = length;
}
return shared_buffer;
}
static PyObject *
set_block_size(PyObject *self, PyObject *args)
{
int new_block_size;
if (PyArg_ParseTuple(args, "i", &new_block_size)) {
if (shared_buffer)
free(shared_buffer);
block_size = new_block_size;
buffer_size = lzo1x_worst_compress(block_size);
shared_buffer = malloc(buffer_size);
omit_headers = 1;
}
Py_INCREF(Py_True);
return Py_True;
}
static PyObject *
lzo_compress(PyObject *self, PyObject *args)
{
const unsigned char *input;
unsigned char *output;
unsigned int inlen, status;
lzo_uint outlen;
/* Get the uncompressed string and its length. */
if (!PyArg_ParseTuple(args, "s#", &input, &inlen))
return NULL;
/* Make sure never to touch unallocated memory. */
if (omit_headers && inlen > block_size)
return PyErr_Format(PyExc_ValueError, "The given input of %i bytes is larger than the configured block size of %i bytes!", block_size, inlen);
/* Allocate the working memory on first use? */
if (!working_memory && !(working_memory = malloc(LZO1X_999_MEM_COMPRESS)))
return PyErr_NoMemory();
/* Allocate the output buffer. */
outlen = lzo1x_worst_compress(inlen);
output = get_buffer(ADD_SIZE(outlen));
if (!output)
return PyErr_NoMemory();
/* Store the input size in the header of the compressed block? */
if (!omit_headers)
*((int*)output) = inlen;
/* Compress the input string. The default LZO compression function is
* lzo1x_1_compress(). There's also variants like lzo1x_1_15_compress() which
* is faster and lzo1x_999_compress() which achieves higher compression. */
status = lzo1x_1_15_compress(input, inlen, ADD_SIZE(output), &outlen, working_memory);
if (status != LZO_E_OK)
return PyErr_Format(PyExc_Exception, "lzo_compress() failed with error code %i!", status);
/* Return the compressed string. */
return Py_BuildValue("s#", output, ADD_SIZE(outlen));
}
static PyObject *
lzo_decompress(PyObject *self, PyObject *args)
{
const unsigned char *input;
unsigned char *output;
int inlen, outlen_expected = 0, status;
lzo_uint outlen_actual;
/* Get the compressed string and its length. */
if (!PyArg_ParseTuple(args, "s#", &input, &inlen))
return NULL;
/* Get the length of the uncompressed string? */
if (!omit_headers)
outlen_expected = *((int*)input);
/* Allocate the output buffer. */
output = get_buffer(outlen_expected);
if (!output)
return PyErr_NoMemory();
/* Decompress the compressed string. */
status = lzo1x_decompress(ADD_SIZE(input), SUB_SIZE(inlen), output, &outlen_actual, NULL);
if (status != LZO_E_OK)
return PyErr_Format(PyExc_Exception, "lzo_decompress() failed with error code %i!", status);
/* Verify the length of the uncompressed data? */
if (!omit_headers && outlen_expected != outlen_actual)
return PyErr_Format(PyExc_Exception, "The expected length (%i) doesn't match the actual uncompressed length (%i)!", outlen_expected, (int)outlen_actual);
/* Return the decompressed string. */
return Py_BuildValue("s#", output, (int)outlen_actual);
}
static PyMethodDef functions[] = {
{ "compress", lzo_compress, METH_VARARGS, "Compress a string using the LZO algorithm." },
{ "decompress", lzo_decompress, METH_VARARGS, "Decompress a string that was previously compressed using the compress() function of this same module." },
{ "set_block_size", set_block_size, METH_VARARGS, "Set the max. length of the strings you will be compressing and/or decompressing so that the LZO module can allocate a single buffer shared by all lzo.compress() and lzo.decompress() calls." },
{ NULL, NULL, 0, NULL }
};
PyMODINIT_FUNC
initlzo(void)
{
int status;
if ((status = lzo_init()) != LZO_E_OK)
PyErr_Format(PyExc_Exception, "Failed to initialize the LZO library! (lzo_init() failed with error code %i)", status);
else if (!Py_InitModule("lzo", functions))
PyErr_Format(PyExc_Exception, "Failed to register module functions!");
}
/* vim: set ts=2 sw=2 et : */