#include "Python.h"

#include <stdio.h>

#include <bzlib.h>

#include "structmember.h"

#ifdef WITH_THREAD

#include "pythread.h"

static char __author__[] =

";

#define BUF(v) PyString_AS_STRING((PyStringObject *)v)

#define MODE_CLOSED 0

#define MODE_READ 1

#define MODE_READ_EOF 2

#define MODE_WRITE 3

#define BZ2FileObject_Check(v) ((v)->ob_type == &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) PyThread_acquire_lock(obj->lock, 1)

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

#define ACQUIRE_LOCK(obj)

#define RELEASE_LOCK(obj)

#ifdef WITH_UNIVERSAL_NEWLINES

#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

PyObject *file;

char* f_buf; /* Allocated readahead buffer */

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

char* f_bufptr; /* Current buffer position */

int f_softspace; /* Flag used by 'print' command */

#ifdef WITH_UNIVERSAL_NEWLINES

int f_univ_newline; /* Handle any newline convention */

int f_newlinetypes; /* Types of newlines seen */

int f_skipnextlf; /* Skip next \n */

BZFILE *fp;

int mode;

long pos;

long 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

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;

#ifdef WITH_UNIVERSAL_NEWLINES

int newlinetypes = f->f_newlinetypes;

int skipnextlf = f->f_skipnextlf;

int univ_newline = f->f_univ_newline;

total_v_size = n > 0 ? n : 100;

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

if (v == NULL)

return NULL;

buf = BUF(v);

end = buf + total_v_size;

for (;;) {

Py_BEGIN_ALLOW_THREADS

#ifdef WITH_UNIVERSAL_NEWLINES

if (univ_newline) {

while (1) {

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

f->pos++;

if (bzerror != BZ_OK || buf == end)

break;

if (skipnextlf) {

skipnextlf = 0;

if (c == '\n') {

/* Seeing a \n here with

newlinetypes |= NEWLINE_CRLF;

BZ2_bzRead(&bzerror, f->fp,

&c, 1);

if (bzerror != BZ_OK)

break;

} else {

newlinetypes |= NEWLINE_CR;

}

}

if (c == '\r') {

skipnextlf = 1;

c = '\n';

} else if ( c == '\n')

newlinetypes |= NEWLINE_LF;

*buf++ = c;

if (c == '\n') break;

}

if (bzerror == BZ_STREAM_END && skipnextlf)

newlinetypes |= NEWLINE_CR;

} else /* If not universal newlines use the normal loop */

do {

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

f->pos++;

*buf++ = c;

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

Py_END_ALLOW_THREADS

#ifdef WITH_UNIVERSAL_NEWLINES

f->f_newlinetypes = newlinetypes;

f->f_skipnextlf = skipnextlf;

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 (_PyString_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)

_PyString_Resize(&v, used_v_size);

return v;

}

#ifndef WITH_UNIVERSAL_NEWLINES

#define Util_UnivNewlineRead(a,b,c,d,e) BZ2_bzRead(a,b,c,d)

Util_UnivNewlineRead(int *bzerror, BZFILE *stream,

char* buf, size_t n, BZ2FileObject *f)

{

char *dst = buf;

int newlinetypes, skipnextlf;

assert(buf != NULL);

assert(stream != NULL);

if (!f->f_univ_newline)

return BZ2_bzRead(bzerror, stream, buf, n);

newlinetypes = f->f_newlinetypes;

skipnextlf = f->f_skipnextlf;

/* Invariant: n is the number of bytes remaining to be filled

while (n) {

size_t nread;

int shortread;

char *src = dst;

nread = BZ2_bzRead(bzerror, stream, dst, n);

assert(nread <= n);

n -= nread; /* assuming 1 byte out for each in; will adjust */

shortread = n != 0; /* true iff EOF or error */

while (nread--) {

char c = *src++;

if (c == '\r') {

/* Save as LF and set flag to skip next LF. */

*dst++ = '\n';

skipnextlf = 1;

}

else if (skipnextlf && c == '\n') {

/* Skip LF, and remember we saw CR LF. */

skipnextlf = 0;

newlinetypes |= NEWLINE_CRLF;

++n;

}

else {

/* Normal char to be stored in buffer. Also

if (c == '\n')

newlinetypes |= NEWLINE_LF;

else if (skipnextlf)

newlinetypes |= NEWLINE_CR;

*dst++ = c;

skipnextlf = 0;

}

}

if (shortread) {

/* If this is EOF, update type flags. */

if (skipnextlf && *bzerror == BZ_STREAM_END)

newlinetypes |= NEWLINE_CR;

break;

}

}

f->f_newlinetypes = newlinetypes;

f->f_skipnextlf = skipnextlf;

return dst - buf;

}

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) {

return -1;

}

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

return -1;

}

Py_BEGIN_ALLOW_THREADS

chunksize = Util_UnivNewlineRead(&bzerror, f->fp, f->f_buf,

bufsize, f);

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 PyStringObject *

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

{

PyStringObject* 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 (PyStringObject *)

PyString_FromStringAndSize(NULL, skip);

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

if (bufptr != NULL) {

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

len = bufptr - f->f_bufptr;

s = (PyStringObject *)

PyString_FromStringAndSize(NULL, skip+len);

if (s == NULL)

return NULL;

memcpy(PyString_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(PyString_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 = PyString_FromString("");

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;

}

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 = PyString_FromStringAndSize((char *)NULL, buffersize);

if (ret == NULL)

goto cleanup;

bytesread = 0;

for (;;) {

Py_BEGIN_ALLOW_THREADS

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

BUF(ret)+bytesread,

buffersize-bytesread,

self);

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 (_PyString_Resize(&ret, buffersize) < 0)

goto cleanup;

} else {

break;

}

}

if (bytesread != buffersize)

_PyString_Resize(&ret, bytesread);

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 = PyString_FromString("");

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;

}

if (sizehint == 0)

ret = PyString_FromString("");

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;

}

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

goto cleanup;

for (;;) {

Py_BEGIN_ALLOW_THREADS

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

buffer+nfilled,

buffersize-nfilled, self);

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 (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 = PyString_FromStringAndSize(

NULL, buffersize);

if (big_buffer == NULL)

goto error;

buffer = PyString_AS_STRING(big_buffer);

memcpy(buffer, small_buffer, nfilled);

}

else {

/* Grow the big buffer */

_PyString_Resize(&big_buffer, buffersize);

buffer = PyString_AS_STRING(big_buffer);

}

continue;

}

end = buffer+nfilled+nread;

q = buffer;

do {

/* Process complete lines */

p++;

line = PyString_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);

} while (p != NULL);

/* 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 = PyString_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;

}

PyString_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_xreadlines__doc__,

");

PyDoc_STRVAR(BZ2File_write__doc__,

");

static PyObject *

BZ2File_write(BZ2FileObject *self, PyObject *args)

{

PyObject *ret = NULL;

char *buf;

int len;

int bzerror;

if (!PyArg_ParseTuple(args, "s#", &buf, &len))

return NULL;

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;;

}

self->f_softspace = 0;

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:

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);

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 strings. If not,

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

PyObject *v = PyList_GET_ITEM(list, i);

if (!PyString_Check(v)) {

const char *buffer;

int len;

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

PyErr_SetString(PyExc_TypeError,

"writelines() "

"argument must be "

"a sequence of "

"strings");

goto error;

}

line = PyString_FromStringAndSize(buffer,

len);

if (line == NULL)

goto error;

Py_DECREF(v);

PyList_SET_ITEM(list, i, line);

}

}

self->f_softspace = 0;

/* 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 = PyString_GET_SIZE(line);

BZ2_bzWrite (&bzerror, self->fp,

PyString_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;

long offset;

char small_buffer[SMALLCHUNK];

char *buffer = small_buffer;

size_t buffersize = SMALLCHUNK;

int bytesread = 0;

int readsize;

int chunksize;

int bzerror;

int rewind = 0;

PyObject *ret = NULL;

if (!PyArg_ParseTuple(args, "l|i:seek", &offset, &where))

return NULL;

ACQUIRE_LOCK(self);

Util_DropReadAhead(self);

switch (self->mode) {

case MODE_READ:

case MODE_READ_EOF:

break;

case MODE_CLOSED: