#include "Python.h"

#include <stdio.h>

#include <bzlib.h>

#include "structmember.h"

#ifdef WITH_THREAD

#include "pythread.h"

static char __author__[] =

";

#if !defined(HAVE_LARGEFILE_SUPPORT)

typedef off_t Py_off_t;

#elif SIZEOF_OFF_T >= 8

typedef off_t Py_off_t;

#elif SIZEOF_FPOS_T >= 8

typedef fpos_t Py_off_t;

#error "Large file support, but neither off_t nor fpos_t is large enough."

#define BUF(v) PyBytes_AS_STRING(v)

#define MODE_CLOSED 0

#define MODE_READ 1

#define MODE_READ_EOF 2

#define MODE_WRITE 3

#define BZ2FileObject_Check(v) (Py_TYPE(v) == &BZ2File_Type)

#ifdef BZ_CONFIG_ERROR

#if SIZEOF_LONG >= 8

#define BZS_TOTAL_OUT(bzs) \

(((long)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)

#elif SIZEOF_LONG_LONG >= 8

#define BZS_TOTAL_OUT(bzs) \

(((PY_LONG_LONG)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)

#define BZS_TOTAL_OUT(bzs) \

bzs->total_out_lo32

#else /* ! BZ_CONFIG_ERROR */

#define BZ2_bzRead bzRead

#define BZ2_bzReadOpen bzReadOpen

#define BZ2_bzReadClose bzReadClose

#define BZ2_bzWrite bzWrite

#define BZ2_bzWriteOpen bzWriteOpen

#define BZ2_bzWriteClose bzWriteClose

#define BZ2_bzCompress bzCompress

#define BZ2_bzCompressInit bzCompressInit

#define BZ2_bzCompressEnd bzCompressEnd

#define BZ2_bzDecompress bzDecompress

#define BZ2_bzDecompressInit bzDecompressInit

#define BZ2_bzDecompressEnd bzDecompressEnd

#define BZS_TOTAL_OUT(bzs) bzs->total_out

#endif /* ! BZ_CONFIG_ERROR */

#ifdef WITH_THREAD

#define ACQUIRE_LOCK(obj) do { \

if (!PyThread_acquire_lock(obj->lock, 0)) { \

Py_BEGIN_ALLOW_THREADS \

PyThread_acquire_lock(obj->lock, 1); \

Py_END_ALLOW_THREADS \

} } while(0)

#define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)

#define ACQUIRE_LOCK(obj)

#define RELEASE_LOCK(obj)

#define NEWLINE_UNKNOWN 0 /* No newline seen, yet */

#define NEWLINE_CR 1 /* \r newline seen */

#define NEWLINE_LF 2 /* \n newline seen */

#define NEWLINE_CRLF 4 /* \r\n newline seen */

typedef struct {

PyObject_HEAD

FILE *rawfp;

char* f_buf; /* Allocated readahead buffer */

char* f_bufend; /* Points after last occupied position */

char* f_bufptr; /* Current buffer position */

BZFILE *fp;

int mode;

Py_off_t pos;

Py_off_t size;

#ifdef WITH_THREAD

PyThread_type_lock lock;

} BZ2FileObject;

typedef struct {

PyObject_HEAD

bz_stream bzs;

int running;

#ifdef WITH_THREAD

PyThread_type_lock lock;

} BZ2CompObject;

typedef struct {

PyObject_HEAD

bz_stream bzs;

int running;

PyObject *unused_data;

#ifdef WITH_THREAD

PyThread_type_lock lock;

} BZ2DecompObject;

static int

check_iterbuffered(BZ2FileObject *f)

{

if (f->f_buf != NULL &&

(f->f_bufend - f->f_bufptr) > 0 &&

f->f_buf[0] != '\0') {

PyErr_SetString(PyExc_ValueError,

"Mixing iteration and read methods would lose data");

return -1;

}

return 0;

}

static int

Util_CatchBZ2Error(int bzerror)

{

int ret = 0;

switch(bzerror) {

case BZ_OK:

case BZ_STREAM_END:

break;

#ifdef BZ_CONFIG_ERROR

case BZ_CONFIG_ERROR:

PyErr_SetString(PyExc_SystemError,

"the bz2 library was not compiled "

"correctly");

ret = 1;

break;

case BZ_PARAM_ERROR:

PyErr_SetString(PyExc_ValueError,

"the bz2 library has received wrong "

"parameters");

ret = 1;

break;

case BZ_MEM_ERROR:

PyErr_NoMemory();

ret = 1;

break;

case BZ_DATA_ERROR:

case BZ_DATA_ERROR_MAGIC:

PyErr_SetString(PyExc_IOError, "invalid data stream");

ret = 1;

break;

case BZ_IO_ERROR:

PyErr_SetString(PyExc_IOError, "unknown IO error");

ret = 1;

break;

case BZ_UNEXPECTED_EOF:

PyErr_SetString(PyExc_EOFError,

"compressed file ended before the "

"logical end-of-stream was detected");

ret = 1;

break;

case BZ_SEQUENCE_ERROR:

PyErr_SetString(PyExc_RuntimeError,

"wrong sequence of bz2 library "

"commands used");

ret = 1;

break;

}

return ret;

}

#if BUFSIZ < 8192

#define SMALLCHUNK 8192

#define SMALLCHUNK BUFSIZ

#if SIZEOF_INT < 4

#define BIGCHUNK (512 * 32)

#define BIGCHUNK (512 * 1024)

static size_t

Util_NewBufferSize(size_t currentsize)

{

if (currentsize > SMALLCHUNK) {

/* Keep doubling until we reach BIGCHUNK;

if (currentsize <= BIGCHUNK)

return currentsize + currentsize;

else

return currentsize + BIGCHUNK;

}

return currentsize + SMALLCHUNK;

}

static PyObject *

Util_GetLine(BZ2FileObject *f, int n)

{

char c;

char *buf, *end;

size_t total_v_size; /* total # of slots in buffer */

size_t used_v_size; /* # used slots in buffer */

size_t increment; /* amount to increment the buffer */

PyObject *v;

int bzerror;

int bytes_read;

total_v_size = n > 0 ? n : 100;

v = PyBytes_FromStringAndSize((char *)NULL, total_v_size);

if (v == NULL)

return NULL;

buf = BUF(v);

end = buf + total_v_size;

for (;;) {

Py_BEGIN_ALLOW_THREADS

do {

bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);

f->pos++;

if (bytes_read == 0)

break;

*buf++ = c;

} while (bzerror == BZ_OK && c != '\n' && buf != end);

Py_END_ALLOW_THREADS

if (bzerror == BZ_STREAM_END) {

f->size = f->pos;

f->mode = MODE_READ_EOF;

break;

} else if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

Py_DECREF(v);

return NULL;

}

if (c == '\n')

break;

/* Must be because buf == end */

if (n > 0)

break;

used_v_size = total_v_size;

increment = total_v_size >> 2; /* mild exponential growth */

total_v_size += increment;

if (total_v_size > INT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"line is longer than a Python string can hold");

Py_DECREF(v);

return NULL;

}

if (_PyBytes_Resize(&v, total_v_size) < 0) {

return NULL;

}

buf = BUF(v) + used_v_size;

end = BUF(v) + total_v_size;

}

used_v_size = buf - BUF(v);

if (used_v_size != total_v_size) {

if (_PyBytes_Resize(&v, used_v_size) < 0) {

v = NULL;

}

}

return v;

}

static void

Util_DropReadAhead(BZ2FileObject *f)

{

if (f->f_buf != NULL) {

PyMem_Free(f->f_buf);

f->f_buf = NULL;

}

}

static int

Util_ReadAhead(BZ2FileObject *f, int bufsize)

{

int chunksize;

int bzerror;

if (f->f_buf != NULL) {

if((f->f_bufend - f->f_bufptr) >= 1)

return 0;

else

Util_DropReadAhead(f);

}

if (f->mode == MODE_READ_EOF) {

f->f_bufptr = f->f_buf;

f->f_bufend = f->f_buf;

return 0;

}

if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {

PyErr_NoMemory();

return -1;

}

Py_BEGIN_ALLOW_THREADS

chunksize = BZ2_bzRead(&bzerror, f->fp, f->f_buf, bufsize);

Py_END_ALLOW_THREADS

f->pos += chunksize;

if (bzerror == BZ_STREAM_END) {

f->size = f->pos;

f->mode = MODE_READ_EOF;

} else if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

Util_DropReadAhead(f);

return -1;

}

f->f_bufptr = f->f_buf;

f->f_bufend = f->f_buf + chunksize;

return 0;

}

static PyBytesObject *

Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)

{

PyBytesObject* s;

char *bufptr;

char *buf;

int len;

if (f->f_buf == NULL)

if (Util_ReadAhead(f, bufsize) < 0)

return NULL;

len = f->f_bufend - f->f_bufptr;

if (len == 0)

return (PyBytesObject *)

PyBytes_FromStringAndSize(NULL, skip);

bufptr = memchr(f->f_bufptr, '\n', len);

if (bufptr != NULL) {

bufptr++; /* Count the '\n' */

len = bufptr - f->f_bufptr;

s = (PyBytesObject *)

PyBytes_FromStringAndSize(NULL, skip+len);

if (s == NULL)

return NULL;

memcpy(PyBytes_AS_STRING(s)+skip, f->f_bufptr, len);

f->f_bufptr = bufptr;

if (bufptr == f->f_bufend)

Util_DropReadAhead(f);

} else {

bufptr = f->f_bufptr;

buf = f->f_buf;

f->f_buf = NULL; /* Force new readahead buffer */

s = Util_ReadAheadGetLineSkip(f, skip+len,

bufsize + (bufsize>>2));

if (s == NULL) {

PyMem_Free(buf);

return NULL;

}

memcpy(PyBytes_AS_STRING(s)+skip, bufptr, len);

PyMem_Free(buf);

}

return s;

}

PyDoc_STRVAR(BZ2File_read__doc__,

");

static PyObject *

BZ2File_read(BZ2FileObject *self, PyObject *args)

{

long bytesrequested = -1;

size_t bytesread, buffersize, chunksize;

int bzerror;

PyObject *ret = NULL;

if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))

return NULL;

ACQUIRE_LOCK(self);

switch (self->mode) {

case MODE_READ:

break;

case MODE_READ_EOF:

ret = PyBytes_FromStringAndSize("", 0);

goto cleanup;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto cleanup;

default:

PyErr_SetString(PyExc_IOError,

"file is not ready for reading");

goto cleanup;

}

/* refuse to mix with f.next() */

if (check_iterbuffered(self))

goto cleanup;

if (bytesrequested < 0)

buffersize = Util_NewBufferSize((size_t)0);

else

buffersize = bytesrequested;

if (buffersize > INT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"requested number of bytes is "

"more than a Python string can hold");

goto cleanup;

}

ret = PyBytes_FromStringAndSize((char *)NULL, buffersize);

if (ret == NULL || buffersize == 0)

goto cleanup;

bytesread = 0;

for (;;) {

Py_BEGIN_ALLOW_THREADS

chunksize = BZ2_bzRead(&bzerror, self->fp,

BUF(ret)+bytesread,

buffersize-bytesread);

self->pos += chunksize;

Py_END_ALLOW_THREADS

bytesread += chunksize;

if (bzerror == BZ_STREAM_END) {

self->size = self->pos;

self->mode = MODE_READ_EOF;

break;

} else if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

Py_DECREF(ret);

ret = NULL;

goto cleanup;

}

if (bytesrequested < 0) {

buffersize = Util_NewBufferSize(buffersize);

if (_PyBytes_Resize(&ret, buffersize) < 0) {

ret = NULL;

goto cleanup;

}

} else {

break;

}

}

if (bytesread != buffersize) {

if (_PyBytes_Resize(&ret, bytesread) < 0) {

ret = NULL;

}

}

cleanup:

RELEASE_LOCK(self);

return ret;

}

PyDoc_STRVAR(BZ2File_readline__doc__,

");

static PyObject *

BZ2File_readline(BZ2FileObject *self, PyObject *args)

{

PyObject *ret = NULL;

int sizehint = -1;

if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))

return NULL;

ACQUIRE_LOCK(self);

switch (self->mode) {

case MODE_READ:

break;

case MODE_READ_EOF:

ret = PyBytes_FromStringAndSize("", 0);

goto cleanup;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto cleanup;

default:

PyErr_SetString(PyExc_IOError,

"file is not ready for reading");

goto cleanup;

}

/* refuse to mix with f.next() */

if (check_iterbuffered(self))

goto cleanup;

if (sizehint == 0)

ret = PyBytes_FromStringAndSize("", 0);

else

ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);

cleanup:

RELEASE_LOCK(self);

return ret;

}

PyDoc_STRVAR(BZ2File_readlines__doc__,

");

static PyObject *

BZ2File_readlines(BZ2FileObject *self, PyObject *args)

{

long sizehint = 0;

PyObject *list = NULL;

PyObject *line;

char small_buffer[SMALLCHUNK];

char *buffer = small_buffer;

size_t buffersize = SMALLCHUNK;

PyObject *big_buffer = NULL;

size_t nfilled = 0;

size_t nread;

size_t totalread = 0;

char *p, *q, *end;

int err;

int shortread = 0;

int bzerror;

if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))

return NULL;

ACQUIRE_LOCK(self);

switch (self->mode) {

case MODE_READ:

break;

case MODE_READ_EOF:

list = PyList_New(0);

goto cleanup;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto cleanup;

default:

PyErr_SetString(PyExc_IOError,

"file is not ready for reading");

goto cleanup;

}

/* refuse to mix with f.next() */

if (check_iterbuffered(self))

goto cleanup;

if ((list = PyList_New(0)) == NULL)

goto cleanup;

for (;;) {

Py_BEGIN_ALLOW_THREADS

nread = BZ2_bzRead(&bzerror, self->fp,

buffer+nfilled, buffersize-nfilled);

self->pos += nread;

Py_END_ALLOW_THREADS

if (bzerror == BZ_STREAM_END) {

self->size = self->pos;

self->mode = MODE_READ_EOF;

if (nread == 0) {

sizehint = 0;

break;

}

shortread = 1;

} else if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

error:

Py_DECREF(list);

list = NULL;

goto cleanup;

}

totalread += nread;

p = memchr(buffer+nfilled, '\n', nread);

if (!shortread && p == NULL) {

/* Need a larger buffer to fit this line */

nfilled += nread;

buffersize *= 2;

if (buffersize > INT_MAX) {

PyErr_SetString(PyExc_OverflowError,

"line is longer than a Python string can hold");

goto error;

}

if (big_buffer == NULL) {

/* Create the big buffer */

big_buffer = PyBytes_FromStringAndSize(

NULL, buffersize);

if (big_buffer == NULL)

goto error;

buffer = PyBytes_AS_STRING(big_buffer);

memcpy(buffer, small_buffer, nfilled);

}

else {

/* Grow the big buffer */

if (_PyBytes_Resize(&big_buffer, buffersize) < 0){

big_buffer = NULL;

goto error;

}

buffer = PyBytes_AS_STRING(big_buffer);

}

continue;

}

end = buffer+nfilled+nread;

q = buffer;

while (p != NULL) {

/* Process complete lines */

p++;

line = PyBytes_FromStringAndSize(q, p-q);

if (line == NULL)

goto error;

err = PyList_Append(list, line);

Py_DECREF(line);

if (err != 0)

goto error;

q = p;

p = memchr(q, '\n', end-q);

}

/* Move the remaining incomplete line to the start */

nfilled = end-q;

memmove(buffer, q, nfilled);

if (sizehint > 0)

if (totalread >= (size_t)sizehint)

break;

if (shortread) {

sizehint = 0;

break;

}

}

if (nfilled != 0) {

/* Partial last line */

line = PyBytes_FromStringAndSize(buffer, nfilled);

if (line == NULL)

goto error;

if (sizehint > 0) {

/* Need to complete the last line */

PyObject *rest = Util_GetLine(self, 0);

if (rest == NULL) {

Py_DECREF(line);

goto error;

}

PyBytes_Concat(&line, rest);

Py_DECREF(rest);

if (line == NULL)

goto error;

}

err = PyList_Append(list, line);

Py_DECREF(line);

if (err != 0)

goto error;

}

cleanup:

RELEASE_LOCK(self);

if (big_buffer) {

Py_DECREF(big_buffer);

}

return list;

}

PyDoc_STRVAR(BZ2File_write__doc__,

");

static PyObject *

BZ2File_write(BZ2FileObject *self, PyObject *args)

{

PyObject *ret = NULL;

Py_buffer pbuf;

char *buf;

int len;

int bzerror;

if (!PyArg_ParseTuple(args, "y*:write", &pbuf))

return NULL;

buf = pbuf.buf;

len = pbuf.len;

ACQUIRE_LOCK(self);

switch (self->mode) {

case MODE_WRITE:

break;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto cleanup;

default:

PyErr_SetString(PyExc_IOError,

"file is not ready for writing");

goto cleanup;

}

Py_BEGIN_ALLOW_THREADS

BZ2_bzWrite (&bzerror, self->fp, buf, len);

self->pos += len;

Py_END_ALLOW_THREADS

if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

goto cleanup;

}

Py_INCREF(Py_None);

ret = Py_None;

cleanup:

PyBuffer_Release(&pbuf);

RELEASE_LOCK(self);

return ret;

}

PyDoc_STRVAR(BZ2File_writelines__doc__,

");

static PyObject *

BZ2File_writelines(BZ2FileObject *self, PyObject *seq)

{

#define CHUNKSIZE 1000

PyObject *list = NULL;

PyObject *iter = NULL;

PyObject *ret = NULL;

PyObject *line;

int i, j, index, len, islist;

int bzerror;

ACQUIRE_LOCK(self);

switch (self->mode) {

case MODE_WRITE:

break;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto error;

default:

PyErr_SetString(PyExc_IOError,

"file is not ready for writing");

goto error;

}

islist = PyList_Check(seq);

if (!islist) {

iter = PyObject_GetIter(seq);

if (iter == NULL) {

PyErr_SetString(PyExc_TypeError,

"writelines() requires an iterable argument");

goto error;

}

list = PyList_New(CHUNKSIZE);

if (list == NULL)

goto error;

}

/* Strategy: slurp CHUNKSIZE lines into a private list,

for (index = 0; ; index += CHUNKSIZE) {

if (islist) {

Py_XDECREF(list);

list = PyList_GetSlice(seq, index, index+CHUNKSIZE);

if (list == NULL)

goto error;

j = PyList_GET_SIZE(list);

}

else {

for (j = 0; j < CHUNKSIZE; j++) {

line = PyIter_Next(iter);

if (line == NULL) {

if (PyErr_Occurred())

goto error;

break;

}

PyList_SetItem(list, j, line);

}

}

if (j == 0)

break;

/* Check that all entries are indeed byte strings. If not,

for (i = 0; i < j; i++) {

PyObject *v = PyList_GET_ITEM(list, i);

if (!PyBytes_Check(v)) {

const char *buffer;

Py_ssize_t len;

if (PyObject_AsCharBuffer(v, &buffer, &len)) {

PyErr_SetString(PyExc_TypeError,

"writelines() "

"argument must be "

"a sequence of "

"bytes objects");

goto error;

}

line = PyBytes_FromStringAndSize(buffer,

len);

if (line == NULL)

goto error;

Py_DECREF(v);

PyList_SET_ITEM(list, i, line);

}

}

/* Since we are releasing the global lock, the

Py_BEGIN_ALLOW_THREADS

for (i = 0; i < j; i++) {

line = PyList_GET_ITEM(list, i);

len = PyBytes_GET_SIZE(line);

BZ2_bzWrite (&bzerror, self->fp,

PyBytes_AS_STRING(line), len);

if (bzerror != BZ_OK) {

Py_BLOCK_THREADS

Util_CatchBZ2Error(bzerror);

goto error;

}

}

Py_END_ALLOW_THREADS

if (j < CHUNKSIZE)

break;

}

Py_INCREF(Py_None);

ret = Py_None;

error:

RELEASE_LOCK(self);

Py_XDECREF(list);

Py_XDECREF(iter);

return ret;

#undef CHUNKSIZE

}

PyDoc_STRVAR(BZ2File_seek__doc__,

");

static PyObject *

BZ2File_seek(BZ2FileObject *self, PyObject *args)

{

int where = 0;

PyObject *offobj;

Py_off_t offset;

char small_buffer[SMALLCHUNK];

char *buffer = small_buffer;

size_t buffersize = SMALLCHUNK;

Py_off_t bytesread = 0;

size_t readsize;

int chunksize;

int bzerror;

PyObject *ret = NULL;

if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))

return NULL;

#if !defined(HAVE_LARGEFILE_SUPPORT)

offset = PyLong_AsLong(offobj);

offset = PyLong_Check(offobj) ?

PyLong_AsLongLong(offobj) : PyLong_AsLong(offobj);

if (PyErr_Occurred())

return NULL;

ACQUIRE_LOCK(self);

Util_DropReadAhead(self);

switch (self->mode) {

case MODE_READ:

case MODE_READ_EOF:

break;

case MODE_CLOSED:

PyErr_SetString(PyExc_ValueError,

"I/O operation on closed file");

goto cleanup;

default:

PyErr_SetString(PyExc_IOError,

"seek works only while reading");

goto cleanup;

}

if (where == 2) {

if (self->size == -1) {

assert(self->mode != MODE_READ_EOF);

for (;;) {

Py_BEGIN_ALLOW_THREADS

chunksize = BZ2_bzRead(&bzerror, self->fp,

buffer, buffersize);

self->pos += chunksize;

Py_END_ALLOW_THREADS

bytesread += chunksize;

if (bzerror == BZ_STREAM_END) {

break;

} else if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

goto cleanup;

}

}

self->mode = MODE_READ_EOF;

self->size = self->pos;

bytesread = 0;

}

offset = self->size + offset;

} else if (where == 1) {

offset = self->pos + offset;

}

/* Before getting here, offset must be the absolute position the file

if (offset >= self->pos) {

/* we can move forward */

offset -= self->pos;

} else {

/* we cannot move back, so rewind the stream */

BZ2_bzReadClose(&bzerror, self->fp);

if (bzerror != BZ_OK) {

Util_CatchBZ2Error(bzerror);

goto cleanup;

}