Modules/bz2module.c

   1 /*
   2
   3 python-bz2 - python bz2 library interface
   4
   5 Copyright (c) 2002  Gustavo Niemeyer <niemeyer@conectiva.com>
   6 Copyright (c) 2002  Python Software Foundation; All Rights Reserved
   7
   8 */
   9
  10 #include "Python.h"
  11 #include <stdio.h>
  12 #include <bzlib.h>
  13 #include "structmember.h"
  14
  15 #ifdef WITH_THREAD
  16 #include "pythread.h"
  17 #endif
  18
  19 static char __author__[] =
  20 "The bz2 python module was written by:\n\
  21 \n\
  22     Gustavo Niemeyer <niemeyer@conectiva.com>\n\
  23 ";
  24
  25 /* Our very own off_t-like type, 64-bit if possible */
  26 /* copied from Objects/fileobject.c */
  27 #if !defined(HAVE_LARGEFILE_SUPPORT)
  28 typedef off_t Py_off_t;
  29 #elif SIZEOF_OFF_T >= 8
  30 typedef off_t Py_off_t;
  31 #elif SIZEOF_FPOS_T >= 8
  32 typedef fpos_t Py_off_t;
  33 #else
  34 #error "Large file support, but neither off_t nor fpos_t is large enough."
  35 #endif
  36
  37 #define BUF(v) PyBytes_AS_STRING(v)
  38
  39 #define MODE_CLOSED   0
  40 #define MODE_READ     1
  41 #define MODE_READ_EOF 2
  42 #define MODE_WRITE    3
  43
  44 #define BZ2FileObject_Check(v)  (Py_TYPE(v) == &BZ2File_Type)
  45
  46
  47 #ifdef BZ_CONFIG_ERROR
  48
  49 #if SIZEOF_LONG >= 8
  50 #define BZS_TOTAL_OUT(bzs) \
  51         (((long)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)
  52 #elif SIZEOF_LONG_LONG >= 8
  53 #define BZS_TOTAL_OUT(bzs) \
  54         (((PY_LONG_LONG)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)
  55 #else
  56 #define BZS_TOTAL_OUT(bzs) \
  57         bzs->total_out_lo32
  58 #endif
  59
  60 #else /* ! BZ_CONFIG_ERROR */
  61
  62 #define BZ2_bzRead bzRead
  63 #define BZ2_bzReadOpen bzReadOpen
  64 #define BZ2_bzReadClose bzReadClose
  65 #define BZ2_bzWrite bzWrite
  66 #define BZ2_bzWriteOpen bzWriteOpen
  67 #define BZ2_bzWriteClose bzWriteClose
  68 #define BZ2_bzCompress bzCompress
  69 #define BZ2_bzCompressInit bzCompressInit
  70 #define BZ2_bzCompressEnd bzCompressEnd
  71 #define BZ2_bzDecompress bzDecompress
  72 #define BZ2_bzDecompressInit bzDecompressInit
  73 #define BZ2_bzDecompressEnd bzDecompressEnd
  74
  75 #define BZS_TOTAL_OUT(bzs) bzs->total_out
  76
  77 #endif /* ! BZ_CONFIG_ERROR */
  78
  79
  80 #ifdef WITH_THREAD
  81 #define ACQUIRE_LOCK(obj) do { \
  82         if (!PyThread_acquire_lock(obj->lock, 0)) { \
  83                 Py_BEGIN_ALLOW_THREADS \
  84                 PyThread_acquire_lock(obj->lock, 1); \
  85                 Py_END_ALLOW_THREADS \
  86         } } while(0)
  87 #define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)
  88 #else
  89 #define ACQUIRE_LOCK(obj)
  90 #define RELEASE_LOCK(obj)
  91 #endif
  92
  93 /* Bits in f_newlinetypes */
  94 #define NEWLINE_UNKNOWN 0       /* No newline seen, yet */
  95 #define NEWLINE_CR 1            /* \r newline seen */
  96 #define NEWLINE_LF 2            /* \n newline seen */
  97 #define NEWLINE_CRLF 4          /* \r\n newline seen */
  98
  99 /* ===================================================================== */
 100 /* Structure definitions. */
 101
 102 typedef struct {
 103         PyObject_HEAD
 104         FILE *rawfp;
 105
 106         char* f_buf;            /* Allocated readahead buffer */
 107         char* f_bufend;         /* Points after last occupied position */
 108         char* f_bufptr;         /* Current buffer position */
 109
 110         BZFILE *fp;
 111         int mode;
 112         Py_off_t pos;
 113         Py_off_t size;
 114 #ifdef WITH_THREAD
 115         PyThread_type_lock lock;
 116 #endif
 117 } BZ2FileObject;
 118
 119 typedef struct {
 120         PyObject_HEAD
 121         bz_stream bzs;
 122         int running;
 123 #ifdef WITH_THREAD
 124         PyThread_type_lock lock;
 125 #endif
 126 } BZ2CompObject;
 127
 128 typedef struct {
 129         PyObject_HEAD
 130         bz_stream bzs;
 131         int running;
 132         PyObject *unused_data;
 133 #ifdef WITH_THREAD
 134         PyThread_type_lock lock;
 135 #endif
 136 } BZ2DecompObject;
 137
 138 /* ===================================================================== */
 139 /* Utility functions. */
 140
 141 static int
 142 Util_CatchBZ2Error(int bzerror)
 143 {
 144         int ret = 0;
 145         switch(bzerror) {
 146                 case BZ_OK:
 147                 case BZ_STREAM_END:
 148                         break;
 149
 150 #ifdef BZ_CONFIG_ERROR
 151                 case BZ_CONFIG_ERROR:
 152                         PyErr_SetString(PyExc_SystemError,
 153                                         "the bz2 library was not compiled "
 154                                         "correctly");
 155                         ret = 1;
 156                         break;
 157 #endif
 158
 159                 case BZ_PARAM_ERROR:
 160                         PyErr_SetString(PyExc_ValueError,
 161                                         "the bz2 library has received wrong "
 162                                         "parameters");
 163                         ret = 1;
 164                         break;
 165
 166                 case BZ_MEM_ERROR:
 167                         PyErr_NoMemory();
 168                         ret = 1;
 169                         break;
 170
 171                 case BZ_DATA_ERROR:
 172                 case BZ_DATA_ERROR_MAGIC:
 173                         PyErr_SetString(PyExc_IOError, "invalid data stream");
 174                         ret = 1;
 175                         break;
 176
 177                 case BZ_IO_ERROR:
 178                         PyErr_SetString(PyExc_IOError, "unknown IO error");
 179                         ret = 1;
 180                         break;
 181
 182                 case BZ_UNEXPECTED_EOF:
 183                         PyErr_SetString(PyExc_EOFError,
 184                                         "compressed file ended before the "
 185                                         "logical end-of-stream was detected");
 186                         ret = 1;
 187                         break;
 188
 189                 case BZ_SEQUENCE_ERROR:
 190                         PyErr_SetString(PyExc_RuntimeError,
 191                                         "wrong sequence of bz2 library "
 192                                         "commands used");
 193                         ret = 1;
 194                         break;
 195         }
 196         return ret;
 197 }
 198
 199 #if BUFSIZ < 8192
 200 #define SMALLCHUNK 8192
 201 #else
 202 #define SMALLCHUNK BUFSIZ
 203 #endif
 204
 205 #if SIZEOF_INT < 4
 206 #define BIGCHUNK  (512 * 32)
 207 #else
 208 #define BIGCHUNK  (512 * 1024)
 209 #endif
 210
 211 /* This is a hacked version of Python's fileobject.c:new_buffersize(). */
 212 static size_t
 213 Util_NewBufferSize(size_t currentsize)
 214 {
 215         if (currentsize > SMALLCHUNK) {
 216                 /* Keep doubling until we reach BIGCHUNK;
 217                    then keep adding BIGCHUNK. */
 218                 if (currentsize <= BIGCHUNK)
 219                         return currentsize + currentsize;
 220                 else
 221                         return currentsize + BIGCHUNK;
 222         }
 223         return currentsize + SMALLCHUNK;
 224 }
 225
 226 /* This is a hacked version of Python's fileobject.c:get_line(). */
 227 static PyObject *
 228 Util_GetLine(BZ2FileObject *f, int n)
 229 {
 230         char c;
 231         char *buf, *end;
 232         size_t total_v_size;    /* total # of slots in buffer */
 233         size_t used_v_size;     /* # used slots in buffer */
 234         size_t increment;       /* amount to increment the buffer */
 235         PyObject *v;
 236         int bzerror;
 237         int bytes_read;
 238
 239         total_v_size = n > 0 ? n : 100;
 240         v = PyBytes_FromStringAndSize((char *)NULL, total_v_size);
 241         if (v == NULL)
 242                 return NULL;
 243
 244         buf = BUF(v);
 245         end = buf + total_v_size;
 246
 247         for (;;) {
 248                 Py_BEGIN_ALLOW_THREADS
 249                 do {
 250                         bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
 251                         f->pos++;
 252                         if (bytes_read == 0)
 253                                 break;
 254                         *buf++ = c;
 255                 } while (bzerror == BZ_OK && c != '\n' && buf != end);
 256                 Py_END_ALLOW_THREADS
 257                 if (bzerror == BZ_STREAM_END) {
 258                         f->size = f->pos;
 259                         f->mode = MODE_READ_EOF;
 260                         break;
 261                 } else if (bzerror != BZ_OK) {
 262                         Util_CatchBZ2Error(bzerror);
 263                         Py_DECREF(v);
 264                         return NULL;
 265                 }
 266                 if (c == '\n')
 267                         break;
 268                 /* Must be because buf == end */
 269                 if (n > 0)
 270                         break;
 271                 used_v_size = total_v_size;
 272                 increment = total_v_size >> 2; /* mild exponential growth */
 273                 total_v_size += increment;
 274                 if (total_v_size > INT_MAX) {
 275                         PyErr_SetString(PyExc_OverflowError,
 276                             "line is longer than a Python string can hold");
 277                         Py_DECREF(v);
 278                         return NULL;
 279                 }
 280                 if (_PyBytes_Resize(&v, total_v_size) < 0) {
 281                         return NULL;
 282                 }
 283                 buf = BUF(v) + used_v_size;
 284                 end = BUF(v) + total_v_size;
 285         }
 286
 287         used_v_size = buf - BUF(v);
 288         if (used_v_size != total_v_size) {
 289                 if (_PyBytes_Resize(&v, used_v_size) < 0) {
 290                         v = NULL;
 291                 }
 292         }
 293         return v;
 294 }
 295
 296 /* This is a hacked version of Python's fileobject.c:drop_readahead(). */
 297 static void
 298 Util_DropReadAhead(BZ2FileObject *f)
 299 {
 300         if (f->f_buf != NULL) {
 301                 PyMem_Free(f->f_buf);
 302                 f->f_buf = NULL;
 303         }
 304 }
 305
 306 /* This is a hacked version of Python's fileobject.c:readahead(). */
 307 static int
 308 Util_ReadAhead(BZ2FileObject *f, int bufsize)
 309 {
 310         int chunksize;
 311         int bzerror;
 312
 313         if (f->f_buf != NULL) {
 314                 if((f->f_bufend - f->f_bufptr) >= 1)
 315                         return 0;
 316                 else
 317                         Util_DropReadAhead(f);
 318         }
 319         if (f->mode == MODE_READ_EOF) {
 320                 f->f_bufptr = f->f_buf;
 321                 f->f_bufend = f->f_buf;
 322                 return 0;
 323         }
 324         if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {
 325                 PyErr_NoMemory();
 326                 return -1;
 327         }
 328         Py_BEGIN_ALLOW_THREADS
 329         chunksize = BZ2_bzRead(&bzerror, f->fp, f->f_buf, bufsize);
 330         Py_END_ALLOW_THREADS
 331         f->pos += chunksize;
 332         if (bzerror == BZ_STREAM_END) {
 333                 f->size = f->pos;
 334                 f->mode = MODE_READ_EOF;
 335         } else if (bzerror != BZ_OK) {
 336                 Util_CatchBZ2Error(bzerror);
 337                 Util_DropReadAhead(f);
 338                 return -1;
 339         }
 340         f->f_bufptr = f->f_buf;
 341         f->f_bufend = f->f_buf + chunksize;
 342         return 0;
 343 }
 344
 345 /* This is a hacked version of Python's
 346  * fileobject.c:readahead_get_line_skip(). */
 347 static PyBytesObject *
 348 Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)
 349 {
 350         PyBytesObject* s;
 351         char *bufptr;
 352         char *buf;
 353         int len;
 354
 355         if (f->f_buf == NULL)
 356                 if (Util_ReadAhead(f, bufsize) < 0)
 357                         return NULL;
 358
 359         len = f->f_bufend - f->f_bufptr;
 360         if (len == 0)
 361                 return (PyBytesObject *)
 362                         PyBytes_FromStringAndSize(NULL, skip);
 363         bufptr = memchr(f->f_bufptr, '\n', len);
 364         if (bufptr != NULL) {
 365                 bufptr++;                       /* Count the '\n' */
 366                 len = bufptr - f->f_bufptr;
 367                 s = (PyBytesObject *)
 368                         PyBytes_FromStringAndSize(NULL, skip+len);
 369                 if (s == NULL)
 370                         return NULL;
 371                 memcpy(PyBytes_AS_STRING(s)+skip, f->f_bufptr, len);
 372                 f->f_bufptr = bufptr;
 373                 if (bufptr == f->f_bufend)
 374                         Util_DropReadAhead(f);
 375         } else {
 376                 bufptr = f->f_bufptr;
 377                 buf = f->f_buf;
 378                 f->f_buf = NULL;        /* Force new readahead buffer */
 379                 s = Util_ReadAheadGetLineSkip(f, skip+len,
 380                                               bufsize + (bufsize>>2));
 381                 if (s == NULL) {
 382                         PyMem_Free(buf);
 383                         return NULL;
 384                 }
 385                 memcpy(PyBytes_AS_STRING(s)+skip, bufptr, len);
 386                 PyMem_Free(buf);
 387         }
 388         return s;
 389 }
 390
 391 /* ===================================================================== */
 392 /* Methods of BZ2File. */
 393
 394 PyDoc_STRVAR(BZ2File_read__doc__,
 395 "read([size]) -> string\n\
 396 \n\
 397 Read at most size uncompressed bytes, returned as a string. If the size\n\
 398 argument is negative or omitted, read until EOF is reached.\n\
 399 ");
 400
 401 /* This is a hacked version of Python's fileobject.c:file_read(). */
 402 static PyObject *
 403 BZ2File_read(BZ2FileObject *self, PyObject *args)
 404 {
 405         long bytesrequested = -1;
 406         size_t bytesread, buffersize, chunksize;
 407         int bzerror;
 408         PyObject *ret = NULL;
 409
 410         if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))
 411                 return NULL;
 412
 413         ACQUIRE_LOCK(self);
 414         switch (self->mode) {
 415                 case MODE_READ:
 416                         break;
 417                 case MODE_READ_EOF:
 418                         ret = PyBytes_FromStringAndSize("", 0);
 419                         goto cleanup;
 420                 case MODE_CLOSED:
 421                         PyErr_SetString(PyExc_ValueError,
 422                                         "I/O operation on closed file");
 423                         goto cleanup;
 424                 default:
 425                         PyErr_SetString(PyExc_IOError,
 426                                         "file is not ready for reading");
 427                         goto cleanup;
 428         }
 429
 430         if (bytesrequested < 0)
 431                 buffersize = Util_NewBufferSize((size_t)0);
 432         else
 433                 buffersize = bytesrequested;
 434         if (buffersize > INT_MAX) {
 435                 PyErr_SetString(PyExc_OverflowError,
 436                                 "requested number of bytes is "
 437                                 "more than a Python string can hold");
 438                 goto cleanup;
 439         }
 440         ret = PyBytes_FromStringAndSize((char *)NULL, buffersize);
 441         if (ret == NULL || buffersize == 0)
 442                 goto cleanup;
 443         bytesread = 0;
 444
 445         for (;;) {
 446                 Py_BEGIN_ALLOW_THREADS
 447                 chunksize = BZ2_bzRead(&bzerror, self->fp,
 448                                        BUF(ret)+bytesread,
 449                                        buffersize-bytesread);
 450                 self->pos += chunksize;
 451                 Py_END_ALLOW_THREADS
 452                 bytesread += chunksize;
 453                 if (bzerror == BZ_STREAM_END) {
 454                         self->size = self->pos;
 455                         self->mode = MODE_READ_EOF;
 456                         break;
 457                 } else if (bzerror != BZ_OK) {
 458                         Util_CatchBZ2Error(bzerror);
 459                         Py_DECREF(ret);
 460                         ret = NULL;
 461                         goto cleanup;
 462                 }
 463                 if (bytesrequested < 0) {
 464                         buffersize = Util_NewBufferSize(buffersize);
 465                         if (_PyBytes_Resize(&ret, buffersize) < 0) {
 466                                 ret = NULL;
 467                                 goto cleanup;
 468                         }
 469                 } else {
 470                         break;
 471                 }
 472         }
 473         if (bytesread != buffersize) {
 474                 if (_PyBytes_Resize(&ret, bytesread) < 0) {
 475                         ret = NULL;
 476                 }
 477         }
 478
 479 cleanup:
 480         RELEASE_LOCK(self);
 481         return ret;
 482 }
 483
 484 PyDoc_STRVAR(BZ2File_readline__doc__,
 485 "readline([size]) -> string\n\
 486 \n\
 487 Return the next line from the file, as a string, retaining newline.\n\
 488 A non-negative size argument will limit the maximum number of bytes to\n\
 489 return (an incomplete line may be returned then). Return an empty\n\
 490 string at EOF.\n\
 491 ");
 492
 493 static PyObject *
 494 BZ2File_readline(BZ2FileObject *self, PyObject *args)
 495 {
 496         PyObject *ret = NULL;
 497         int sizehint = -1;
 498
 499         if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))
 500                 return NULL;
 501
 502         ACQUIRE_LOCK(self);
 503         switch (self->mode) {
 504                 case MODE_READ:
 505                         break;
 506                 case MODE_READ_EOF:
 507                         ret = PyBytes_FromStringAndSize("", 0);
 508                         goto cleanup;
 509                 case MODE_CLOSED:
 510                         PyErr_SetString(PyExc_ValueError,
 511                                         "I/O operation on closed file");
 512                         goto cleanup;
 513                 default:
 514                         PyErr_SetString(PyExc_IOError,
 515                                         "file is not ready for reading");
 516                         goto cleanup;
 517         }
 518
 519         if (sizehint == 0)
 520                 ret = PyBytes_FromStringAndSize("", 0);
 521         else
 522                 ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);
 523
 524 cleanup:
 525         RELEASE_LOCK(self);
 526         return ret;
 527 }
 528
 529 PyDoc_STRVAR(BZ2File_readlines__doc__,
 530 "readlines([size]) -> list\n\
 531 \n\
 532 Call readline() repeatedly and return a list of lines read.\n\
 533 The optional size argument, if given, is an approximate bound on the\n\
 534 total number of bytes in the lines returned.\n\
 535 ");
 536
 537 /* This is a hacked version of Python's fileobject.c:file_readlines(). */
 538 static PyObject *
 539 BZ2File_readlines(BZ2FileObject *self, PyObject *args)
 540 {
 541         long sizehint = 0;
 542         PyObject *list = NULL;
 543         PyObject *line;
 544         char small_buffer[SMALLCHUNK];
 545         char *buffer = small_buffer;
 546         size_t buffersize = SMALLCHUNK;
 547         PyObject *big_buffer = NULL;
 548         size_t nfilled = 0;
 549         size_t nread;
 550         size_t totalread = 0;
 551         char *p, *q, *end;
 552         int err;
 553         int shortread = 0;
 554         int bzerror;
 555
 556         if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))
 557                 return NULL;
 558
 559         ACQUIRE_LOCK(self);
 560         switch (self->mode) {
 561                 case MODE_READ:
 562                         break;
 563                 case MODE_READ_EOF:
 564                         list = PyList_New(0);
 565                         goto cleanup;
 566                 case MODE_CLOSED:
 567                         PyErr_SetString(PyExc_ValueError,
 568                                         "I/O operation on closed file");
 569                         goto cleanup;
 570                 default:
 571                         PyErr_SetString(PyExc_IOError,
 572                                         "file is not ready for reading");
 573                         goto cleanup;
 574         }
 575
 576         if ((list = PyList_New(0)) == NULL)
 577                 goto cleanup;
 578
 579         for (;;) {
 580                 Py_BEGIN_ALLOW_THREADS
 581                 nread = BZ2_bzRead(&bzerror, self->fp,
 582                                    buffer+nfilled, buffersize-nfilled);
 583                 self->pos += nread;
 584                 Py_END_ALLOW_THREADS
 585                 if (bzerror == BZ_STREAM_END) {
 586                         self->size = self->pos;
 587                         self->mode = MODE_READ_EOF;
 588                         if (nread == 0) {
 589                                 sizehint = 0;
 590                                 break;
 591                         }
 592                         shortread = 1;
 593                 } else if (bzerror != BZ_OK) {
 594                         Util_CatchBZ2Error(bzerror);
 595                   error:
 596                         Py_DECREF(list);
 597                         list = NULL;
 598                         goto cleanup;
 599                 }
 600                 totalread += nread;
 601                 p = memchr(buffer+nfilled, '\n', nread);
 602                 if (!shortread && p == NULL) {
 603                         /* Need a larger buffer to fit this line */
 604                         nfilled += nread;
 605                         buffersize *= 2;
 606                         if (buffersize > INT_MAX) {
 607                                 PyErr_SetString(PyExc_OverflowError,
 608                                 "line is longer than a Python string can hold");
 609                                 goto error;
 610                         }
 611                         if (big_buffer == NULL) {
 612                                 /* Create the big buffer */
 613                                 big_buffer = PyBytes_FromStringAndSize(
 614                                         NULL, buffersize);
 615                                 if (big_buffer == NULL)
 616                                         goto error;
 617                                 buffer = PyBytes_AS_STRING(big_buffer);
 618                                 memcpy(buffer, small_buffer, nfilled);
 619                         }
 620                         else {
 621                                 /* Grow the big buffer */
 622                                 if (_PyBytes_Resize(&big_buffer, buffersize) < 0){
 623                                         big_buffer = NULL;
 624                                         goto error;
 625                                 }
 626                                 buffer = PyBytes_AS_STRING(big_buffer);
 627                         }
 628                         continue;
 629                 }
 630                 end = buffer+nfilled+nread;
 631                 q = buffer;
 632                 while (p != NULL) {
 633                         /* Process complete lines */
 634                         p++;
 635                         line = PyBytes_FromStringAndSize(q, p-q);
 636                         if (line == NULL)
 637                                 goto error;
 638                         err = PyList_Append(list, line);
 639                         Py_DECREF(line);
 640                         if (err != 0)
 641                                 goto error;
 642                         q = p;
 643                         p = memchr(q, '\n', end-q);
 644                 }
 645                 /* Move the remaining incomplete line to the start */
 646                 nfilled = end-q;
 647                 memmove(buffer, q, nfilled);
 648                 if (sizehint > 0)
 649                         if (totalread >= (size_t)sizehint)
 650                                 break;
 651                 if (shortread) {
 652                         sizehint = 0;
 653                         break;
 654                 }
 655         }
 656         if (nfilled != 0) {
 657                 /* Partial last line */
 658                 line = PyBytes_FromStringAndSize(buffer, nfilled);
 659                 if (line == NULL)
 660                         goto error;
 661                 if (sizehint > 0) {
 662                         /* Need to complete the last line */
 663                         PyObject *rest = Util_GetLine(self, 0);
 664                         if (rest == NULL) {
 665                                 Py_DECREF(line);
 666                                 goto error;
 667                         }
 668                         PyBytes_Concat(&line, rest);
 669                         Py_DECREF(rest);
 670                         if (line == NULL)
 671                                 goto error;
 672                 }
 673                 err = PyList_Append(list, line);
 674                 Py_DECREF(line);
 675                 if (err != 0)
 676                         goto error;
 677         }
 678
 679   cleanup:
 680         RELEASE_LOCK(self);
 681         if (big_buffer) {
 682                 Py_DECREF(big_buffer);
 683         }
 684         return list;
 685 }
 686
 687 PyDoc_STRVAR(BZ2File_write__doc__,
 688 "write(data) -> None\n\
 689 \n\
 690 Write the 'data' string to file. Note that due to buffering, close() may\n\
 691 be needed before the file on disk reflects the data written.\n\
 692 ");
 693
 694 /* This is a hacked version of Python's fileobject.c:file_write(). */
 695 static PyObject *
 696 BZ2File_write(BZ2FileObject *self, PyObject *args)
 697 {
 698         PyObject *ret = NULL;
 699         Py_buffer pbuf;
 700         char *buf;
 701         int len;
 702         int bzerror;
 703
 704         if (!PyArg_ParseTuple(args, "y*:write", &pbuf))
 705                 return NULL;
 706         buf = pbuf.buf;
 707         len = pbuf.len;
 708
 709         ACQUIRE_LOCK(self);
 710         switch (self->mode) {
 711                 case MODE_WRITE:
 712                         break;
 713
 714                 case MODE_CLOSED:
 715                         PyErr_SetString(PyExc_ValueError,
 716                                         "I/O operation on closed file");
 717                         goto cleanup;
 718
 719                 default:
 720                         PyErr_SetString(PyExc_IOError,
 721                                         "file is not ready for writing");
 722                         goto cleanup;
 723         }
 724
 725         Py_BEGIN_ALLOW_THREADS
 726         BZ2_bzWrite (&bzerror, self->fp, buf, len);
 727         self->pos += len;
 728         Py_END_ALLOW_THREADS
 729
 730         if (bzerror != BZ_OK) {
 731                 Util_CatchBZ2Error(bzerror);
 732                 goto cleanup;
 733         }
 734
 735         Py_INCREF(Py_None);
 736         ret = Py_None;
 737
 738 cleanup:
 739         PyBuffer_Release(&pbuf);
 740         RELEASE_LOCK(self);
 741         return ret;
 742 }
 743
 744 PyDoc_STRVAR(BZ2File_writelines__doc__,
 745 "writelines(sequence_of_strings) -> None\n\
 746 \n\
 747 Write the sequence of strings to the file. Note that newlines are not\n\
 748 added. The sequence can be any iterable object producing strings. This is\n\
 749 equivalent to calling write() for each string.\n\
 750 ");
 751
 752 /* This is a hacked version of Python's fileobject.c:file_writelines(). */
 753 static PyObject *
 754 BZ2File_writelines(BZ2FileObject *self, PyObject *seq)
 755 {
 756 #define CHUNKSIZE 1000
 757         PyObject *list = NULL;
 758         PyObject *iter = NULL;
 759         PyObject *ret = NULL;
 760         PyObject *line;
 761         int i, j, index, len, islist;
 762         int bzerror;
 763
 764         ACQUIRE_LOCK(self);
 765         switch (self->mode) {
 766                 case MODE_WRITE:
 767                         break;
 768
 769                 case MODE_CLOSED:
 770                         PyErr_SetString(PyExc_ValueError,
 771                                         "I/O operation on closed file");
 772                         goto error;
 773
 774                 default:
 775                         PyErr_SetString(PyExc_IOError,
 776                                         "file is not ready for writing");
 777                         goto error;
 778         }
 779
 780         islist = PyList_Check(seq);
 781         if  (!islist) {
 782                 iter = PyObject_GetIter(seq);
 783                 if (iter == NULL) {
 784                         PyErr_SetString(PyExc_TypeError,
 785                                 "writelines() requires an iterable argument");
 786                         goto error;
 787                 }
 788                 list = PyList_New(CHUNKSIZE);
 789                 if (list == NULL)
 790                         goto error;
 791         }
 792
 793         /* Strategy: slurp CHUNKSIZE lines into a private list,
 794            checking that they are all strings, then write that list
 795            without holding the interpreter lock, then come back for more. */
 796         for (index = 0; ; index += CHUNKSIZE) {
 797                 if (islist) {
 798                         Py_XDECREF(list);
 799                         list = PyList_GetSlice(seq, index, index+CHUNKSIZE);
 800                         if (list == NULL)
 801                                 goto error;
 802                         j = PyList_GET_SIZE(list);
 803                 }
 804                 else {
 805                         for (j = 0; j < CHUNKSIZE; j++) {
 806                                 line = PyIter_Next(iter);
 807                                 if (line == NULL) {
 808                                         if (PyErr_Occurred())
 809                                                 goto error;
 810                                         break;
 811                                 }
 812                                 PyList_SetItem(list, j, line);
 813                         }
 814                 }
 815                 if (j == 0)
 816                         break;
 817
 818                 /* Check that all entries are indeed byte strings. If not,
 819                    apply the same rules as for file.write() and
 820                    convert the rets to strings. This is slow, but
 821                    seems to be the only way since all conversion APIs
 822                    could potentially execute Python code. */
 823                 for (i = 0; i < j; i++) {
 824                         PyObject *v = PyList_GET_ITEM(list, i);
 825                         if (!PyBytes_Check(v)) {
 826                                 const char *buffer;
 827                                 Py_ssize_t len;
 828                                 if (PyObject_AsCharBuffer(v, &buffer, &len)) {
 829                                         PyErr_SetString(PyExc_TypeError,
 830                                                         "writelines() "
 831                                                         "argument must be "
 832                                                         "a sequence of "
 833                                                         "bytes objects");
 834                                         goto error;
 835                                 }
 836                                 line = PyBytes_FromStringAndSize(buffer,
 837                                                                   len);
 838                                 if (line == NULL)
 839                                         goto error;
 840                                 Py_DECREF(v);
 841                                 PyList_SET_ITEM(list, i, line);
 842                         }
 843                 }
 844
 845                 /* Since we are releasing the global lock, the
 846                    following code may *not* execute Python code. */
 847                 Py_BEGIN_ALLOW_THREADS
 848                 for (i = 0; i < j; i++) {
 849                         line = PyList_GET_ITEM(list, i);
 850                         len = PyBytes_GET_SIZE(line);
 851                         BZ2_bzWrite (&bzerror, self->fp,
 852                                      PyBytes_AS_STRING(line), len);
 853                         if (bzerror != BZ_OK) {
 854                                 Py_BLOCK_THREADS
 855                                 Util_CatchBZ2Error(bzerror);
 856                                 goto error;
 857                         }
 858                 }
 859                 Py_END_ALLOW_THREADS
 860
 861                 if (j < CHUNKSIZE)
 862                         break;
 863         }
 864
 865         Py_INCREF(Py_None);
 866         ret = Py_None;
 867
 868   error:
 869         RELEASE_LOCK(self);
 870         Py_XDECREF(list);
 871         Py_XDECREF(iter);
 872         return ret;
 873 #undef CHUNKSIZE
 874 }
 875
 876 PyDoc_STRVAR(BZ2File_seek__doc__,
 877 "seek(offset [, whence]) -> None\n\
 878 \n\
 879 Move to new file position. Argument offset is a byte count. Optional\n\
 880 argument whence defaults to 0 (offset from start of file, offset\n\
 881 should be >= 0); other values are 1 (move relative to current position,\n\
 882 positive or negative), and 2 (move relative to end of file, usually\n\
 883 negative, although many platforms allow seeking beyond the end of a file).\n\
 884 \n\
 885 Note that seeking of bz2 files is emulated, and depending on the parameters\n\
 886 the operation may be extremely slow.\n\
 887 ");
 888
 889 static PyObject *
 890 BZ2File_seek(BZ2FileObject *self, PyObject *args)
 891 {
 892         int where = 0;
 893         PyObject *offobj;
 894         Py_off_t offset;
 895         char small_buffer[SMALLCHUNK];
 896         char *buffer = small_buffer;
 897         size_t buffersize = SMALLCHUNK;
 898         Py_off_t bytesread = 0;
 899         size_t readsize;
 900         int chunksize;
 901         int bzerror;
 902         PyObject *ret = NULL;
 903
 904         if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))
 905                 return NULL;
 906 #if !defined(HAVE_LARGEFILE_SUPPORT)
 907         offset = PyLong_AsLong(offobj);
 908 #else
 909         offset = PyLong_Check(offobj) ?
 910                 PyLong_AsLongLong(offobj) : PyLong_AsLong(offobj);
 911 #endif
 912         if (PyErr_Occurred())
 913                 return NULL;
 914
 915         ACQUIRE_LOCK(self);
 916         Util_DropReadAhead(self);
 917         switch (self->mode) {
 918                 case MODE_READ:
 919                 case MODE_READ_EOF:
 920                         break;
 921
 922                 case MODE_CLOSED:
 923                         PyErr_SetString(PyExc_ValueError,
 924                                         "I/O operation on closed file");
 925                         goto cleanup;
 926
 927                 default:
 928                         PyErr_SetString(PyExc_IOError,
 929                                         "seek works only while reading");
 930                         goto cleanup;
 931         }
 932
 933         if (where == 2) {
 934                 if (self->size == -1) {
 935                         assert(self->mode != MODE_READ_EOF);
 936                         for (;;) {
 937                                 Py_BEGIN_ALLOW_THREADS
 938                                 chunksize = BZ2_bzRead(&bzerror, self->fp,
 939                                                        buffer, buffersize);
 940                                 self->pos += chunksize;
 941                                 Py_END_ALLOW_THREADS
 942
 943                                 bytesread += chunksize;
 944                                 if (bzerror == BZ_STREAM_END) {
 945                                         break;
 946                                 } else if (bzerror != BZ_OK) {
 947                                         Util_CatchBZ2Error(bzerror);
 948                                         goto cleanup;
 949                                 }
 950                         }
 951                         self->mode = MODE_READ_EOF;
 952                         self->size = self->pos;
 953                         bytesread = 0;
 954                 }
 955                 offset = self->size + offset;
 956         } else if (where == 1) {
 957                 offset = self->pos + offset;
 958         }
 959
 960         /* Before getting here, offset must be the absolute position the file
 961          * pointer should be set to. */
 962
 963         if (offset >= self->pos) {
 964                 /* we can move forward */
 965                 offset -= self->pos;
 966         } else {
 967                 /* we cannot move back, so rewind the stream */
 968                 BZ2_bzReadClose(&bzerror, self->fp);
 969                 if (bzerror != BZ_OK) {
 970                         Util_CatchBZ2Error(bzerror);
 971                         goto cleanup;
 972                 }
 973                 rewind(self->rawfp);
 974                 self->pos = 0;
 975                 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
 976                                           0, 0, NULL, 0);
 977                 if (bzerror != BZ_OK) {
 978                         Util_CatchBZ2Error(bzerror);
 979                         goto cleanup;
 980                 }
 981                 self->mode = MODE_READ;
 982         }
 983
 984         if (offset <= 0 || self->mode == MODE_READ_EOF)
 985                 goto exit;
 986
 987         /* Before getting here, offset must be set to the number of bytes
 988          * to walk forward. */
 989         for (;;) {
 990                 if (offset-bytesread > buffersize)
 991                         readsize = buffersize;
 992                 else
 993                         /* offset might be wider that readsize, but the result
 994                          * of the subtraction is bound by buffersize (see the
 995                          * condition above). buffersize is 8192. */
 996                         readsize = (size_t)(offset-bytesread);
 997                 Py_BEGIN_ALLOW_THREADS
 998                 chunksize = BZ2_bzRead(&bzerror, self->fp, buffer, readsize);
 999                 self->pos += chunksize;
1000                 Py_END_ALLOW_THREADS
1001                 bytesread += chunksize;
1002                 if (bzerror == BZ_STREAM_END) {
1003                         self->size = self->pos;
1004                         self->mode = MODE_READ_EOF;
1005                         break;
1006                 } else if (bzerror != BZ_OK) {
1007                         Util_CatchBZ2Error(bzerror);
1008                         goto cleanup;
1009                 }
1010                 if (bytesread == offset)
1011                         break;
1012         }
1013
1014 exit:
1015         Py_INCREF(Py_None);
1016         ret = Py_None;
1017
1018 cleanup:
1019         RELEASE_LOCK(self);
1020         return ret;
1021 }
1022
1023 PyDoc_STRVAR(BZ2File_tell__doc__,
1024 "tell() -> int\n\
1025 \n\
1026 Return the current file position, an integer (may be a long integer).\n\
1027 ");
1028
1029 static PyObject *
1030 BZ2File_tell(BZ2FileObject *self, PyObject *args)
1031 {
1032         PyObject *ret = NULL;
1033
1034         if (self->mode == MODE_CLOSED) {
1035                 PyErr_SetString(PyExc_ValueError,
1036                                 "I/O operation on closed file");
1037                 goto cleanup;
1038         }
1039
1040 #if !defined(HAVE_LARGEFILE_SUPPORT)
1041         ret = PyLong_FromLong(self->pos);
1042 #else
1043         ret = PyLong_FromLongLong(self->pos);
1044 #endif
1045
1046 cleanup:
1047         return ret;
1048 }
1049
1050 PyDoc_STRVAR(BZ2File_close__doc__,
1051 "close() -> None or (perhaps) an integer\n\
1052 \n\
1053 Close the file. Sets data attribute .closed to true. A closed file\n\
1054 cannot be used for further I/O operations. close() may be called more\n\
1055 than once without error.\n\
1056 ");
1057
1058 static PyObject *
1059 BZ2File_close(BZ2FileObject *self)
1060 {
1061         PyObject *ret = NULL;
1062         int bzerror = BZ_OK;
1063
1064         if (self->mode == MODE_CLOSED) {
1065                 Py_RETURN_NONE;
1066         }
1067
1068         ACQUIRE_LOCK(self);
1069         switch (self->mode) {
1070                 case MODE_READ:
1071                 case MODE_READ_EOF:
1072                         BZ2_bzReadClose(&bzerror, self->fp);
1073                         break;
1074                 case MODE_WRITE:
1075                         BZ2_bzWriteClose(&bzerror, self->fp,
1076                                          0, NULL, NULL);
1077                         break;
1078         }
1079         self->mode = MODE_CLOSED;
1080         fclose(self->rawfp);
1081         self->rawfp = NULL;
1082         if (bzerror == BZ_OK) {
1083                 Py_INCREF(Py_None);
1084                 ret = Py_None;
1085         }
1086         else {
1087                 Util_CatchBZ2Error(bzerror);
1088         }
1089
1090         RELEASE_LOCK(self);
1091         return ret;
1092 }
1093
1094 PyDoc_STRVAR(BZ2File_enter_doc,
1095 "__enter__() -> self.");
1096
1097 static PyObject *
1098 BZ2File_enter(BZ2FileObject *self)
1099 {
1100         if (self->mode == MODE_CLOSED) {
1101                 PyErr_SetString(PyExc_ValueError,
1102                         "I/O operation on closed file");
1103                 return NULL;
1104         }
1105         Py_INCREF(self);
1106         return (PyObject *) self;
1107 }
1108
1109 PyDoc_STRVAR(BZ2File_exit_doc,
1110 "__exit__(*excinfo) -> None.  Closes the file.");
1111
1112 static PyObject *
1113 BZ2File_exit(BZ2FileObject *self, PyObject *args)
1114 {
1115         PyObject *ret = PyObject_CallMethod((PyObject *) self, "close", NULL);
1116         if (!ret)
1117                 /* If error occurred, pass through */
1118                 return NULL;
1119         Py_DECREF(ret);
1120         Py_RETURN_NONE;
1121 }
1122
1123
1124 static PyObject *BZ2File_getiter(BZ2FileObject *self);
1125
1126 static PyMethodDef BZ2File_methods[] = {
1127         {"read", (PyCFunction)BZ2File_read, METH_VARARGS, BZ2File_read__doc__},
1128         {"readline", (PyCFunction)BZ2File_readline, METH_VARARGS, BZ2File_readline__doc__},
1129         {"readlines", (PyCFunction)BZ2File_readlines, METH_VARARGS, BZ2File_readlines__doc__},
1130         {"write", (PyCFunction)BZ2File_write, METH_VARARGS, BZ2File_write__doc__},
1131         {"writelines", (PyCFunction)BZ2File_writelines, METH_O, BZ2File_writelines__doc__},
1132         {"seek", (PyCFunction)BZ2File_seek, METH_VARARGS, BZ2File_seek__doc__},
1133         {"tell", (PyCFunction)BZ2File_tell, METH_NOARGS, BZ2File_tell__doc__},
1134         {"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
1135         {"__enter__", (PyCFunction)BZ2File_enter, METH_NOARGS, BZ2File_enter_doc},
1136         {"__exit__", (PyCFunction)BZ2File_exit, METH_VARARGS, BZ2File_exit_doc},
1137         {NULL,          NULL}           /* sentinel */
1138 };
1139
1140
1141 /* ===================================================================== */
1142 /* Getters and setters of BZ2File. */
1143
1144 static PyObject *
1145 BZ2File_get_closed(BZ2FileObject *self, void *closure)
1146 {
1147         return PyLong_FromLong(self->mode == MODE_CLOSED);
1148 }
1149
1150 static PyGetSetDef BZ2File_getset[] = {
1151         {"closed", (getter)BZ2File_get_closed, NULL,
1152                         "True if the file is closed"},
1153         {NULL}  /* Sentinel */
1154 };
1155
1156
1157 /* ===================================================================== */
1158 /* Slot definitions for BZ2File_Type. */
1159
1160 static int
1161 BZ2File_init(BZ2FileObject *self, PyObject *args, PyObject *kwargs)
1162 {
1163         static char *kwlist[] = {"filename", "mode", "buffering",
1164                                  "compresslevel", 0};
1165         PyObject *name_obj = NULL;
1166         char *name;
1167         char *mode = "r";
1168         int buffering = -1;
1169         int compresslevel = 9;
1170         int bzerror;
1171         int mode_char = 0;
1172
1173         self->size = -1;
1174
1175         if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&|sii:BZ2File",
1176                                          kwlist, PyUnicode_FSConverter, &name_obj,
1177                                          &mode, &buffering,
1178                                          &compresslevel))
1179                 return -1;
1180
1181         if (PyBytes_Check(name_obj))
1182                 name = PyBytes_AsString(name_obj);
1183         else
1184                 name = PyByteArray_AsString(name_obj);
1185         if (compresslevel < 1 || compresslevel > 9) {
1186                 PyErr_SetString(PyExc_ValueError,
1187                                 "compresslevel must be between 1 and 9");
1188                 Py_DECREF(name_obj);
1189                 return -1;
1190         }
1191
1192         for (;;) {
1193                 int error = 0;
1194                 switch (*mode) {
1195                         case 'r':
1196                         case 'w':
1197                                 if (mode_char)
1198                                         error = 1;
1199                                 mode_char = *mode;
1200                                 break;
1201
1202                         case 'b':
1203                                 break;
1204
1205                         default:
1206                                 error = 1;
1207                                 break;
1208                 }
1209                 if (error) {
1210                         PyErr_Format(PyExc_ValueError,
1211                                      "invalid mode char %c", *mode);
1212                         Py_DECREF(name_obj);
1213                         return -1;
1214                 }
1215                 mode++;
1216                 if (*mode == '\0')
1217                         break;
1218         }
1219
1220         if (mode_char == 0) {
1221                 mode_char = 'r';
1222         }
1223
1224         mode = (mode_char == 'r') ? "rb" : "wb";
1225
1226         self->rawfp = fopen(name, mode);
1227         Py_DECREF(name_obj);
1228         if (self->rawfp == NULL) {
1229                 PyErr_SetFromErrno(PyExc_IOError);
1230                 return -1;
1231         }
1232         /* XXX Ignore buffering */
1233
1234         /* From now on, we have stuff to dealloc, so jump to error label
1235          * instead of returning */
1236
1237 #ifdef WITH_THREAD
1238         self->lock = PyThread_allocate_lock();
1239         if (!self->lock) {
1240                 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1241                 goto error;
1242         }
1243 #endif
1244
1245         if (mode_char == 'r')
1246                 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
1247                                           0, 0, NULL, 0);
1248         else
1249                 self->fp = BZ2_bzWriteOpen(&bzerror, self->rawfp,
1250                                            compresslevel, 0, 0);
1251
1252         if (bzerror != BZ_OK) {
1253                 Util_CatchBZ2Error(bzerror);
1254                 goto error;
1255         }
1256
1257         self->mode = (mode_char == 'r') ? MODE_READ : MODE_WRITE;
1258
1259         return 0;
1260
1261 error:
1262         fclose(self->rawfp);
1263         self->rawfp = NULL;
1264 #ifdef WITH_THREAD
1265         if (self->lock) {
1266                 PyThread_free_lock(self->lock);
1267                 self->lock = NULL;
1268         }
1269 #endif
1270         return -1;
1271 }
1272
1273 static void
1274 BZ2File_dealloc(BZ2FileObject *self)
1275 {
1276         int bzerror;
1277 #ifdef WITH_THREAD
1278         if (self->lock)
1279                 PyThread_free_lock(self->lock);
1280 #endif
1281         switch (self->mode) {
1282                 case MODE_READ:
1283                 case MODE_READ_EOF:
1284                         BZ2_bzReadClose(&bzerror, self->fp);
1285                         break;
1286                 case MODE_WRITE:
1287                         BZ2_bzWriteClose(&bzerror, self->fp,
1288                                          0, NULL, NULL);
1289                         break;
1290         }
1291         Util_DropReadAhead(self);
1292         if (self->rawfp != NULL)
1293                 fclose(self->rawfp);
1294         Py_TYPE(self)->tp_free((PyObject *)self);
1295 }
1296
1297 /* This is a hacked version of Python's fileobject.c:file_getiter(). */
1298 static PyObject *
1299 BZ2File_getiter(BZ2FileObject *self)
1300 {
1301         if (self->mode == MODE_CLOSED) {
1302                 PyErr_SetString(PyExc_ValueError,
1303                                 "I/O operation on closed file");
1304                 return NULL;
1305         }
1306         Py_INCREF((PyObject*)self);
1307         return (PyObject *)self;
1308 }
1309
1310 /* This is a hacked version of Python's fileobject.c:file_iternext(). */
1311 #define READAHEAD_BUFSIZE 8192
1312 static PyObject *
1313 BZ2File_iternext(BZ2FileObject *self)
1314 {
1315         PyBytesObject* ret;
1316         ACQUIRE_LOCK(self);
1317         if (self->mode == MODE_CLOSED) {
1318                 RELEASE_LOCK(self);
1319                 PyErr_SetString(PyExc_ValueError,
1320                                 "I/O operation on closed file");
1321                 return NULL;
1322         }
1323         ret = Util_ReadAheadGetLineSkip(self, 0, READAHEAD_BUFSIZE);
1324         RELEASE_LOCK(self);
1325         if (ret == NULL || PyBytes_GET_SIZE(ret) == 0) {
1326                 Py_XDECREF(ret);
1327                 return NULL;
1328         }
1329         return (PyObject *)ret;
1330 }
1331
1332 /* ===================================================================== */
1333 /* BZ2File_Type definition. */
1334
1335 PyDoc_VAR(BZ2File__doc__) =
1336 PyDoc_STR(
1337 "BZ2File(name [, mode='r', buffering=0, compresslevel=9]) -> file object\n\
1338 \n\
1339 Open a bz2 file. The mode can be 'r' or 'w', for reading (default) or\n\
1340 writing. When opened for writing, the file will be created if it doesn't\n\
1341 exist, and truncated otherwise. If the buffering argument is given, 0 means\n\
1342 unbuffered, and larger numbers specify the buffer size. If compresslevel\n\
1343 is given, must be a number between 1 and 9.\n\
1344 Data read is always returned in bytes; data written ought to be bytes.\n\
1345 ");
1346
1347 static PyTypeObject BZ2File_Type = {
1348         PyVarObject_HEAD_INIT(NULL, 0)
1349         "bz2.BZ2File",          /*tp_name*/
1350         sizeof(BZ2FileObject),  /*tp_basicsize*/
1351         0,                      /*tp_itemsize*/
1352         (destructor)BZ2File_dealloc, /*tp_dealloc*/
1353         0,                      /*tp_print*/
1354         0,                      /*tp_getattr*/
1355         0,                      /*tp_setattr*/
1356         0,                      /*tp_reserved*/
1357         0,                      /*tp_repr*/
1358         0,                      /*tp_as_number*/
1359         0,                      /*tp_as_sequence*/
1360         0,                      /*tp_as_mapping*/
1361         0,                      /*tp_hash*/
1362         0,                      /*tp_call*/
1363         0,                      /*tp_str*/
1364         PyObject_GenericGetAttr,/*tp_getattro*/
1365         PyObject_GenericSetAttr,/*tp_setattro*/
1366         0,                      /*tp_as_buffer*/
1367         Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1368         BZ2File__doc__,         /*tp_doc*/
1369         0,                      /*tp_traverse*/
1370         0,                      /*tp_clear*/
1371         0,                      /*tp_richcompare*/
1372         0,                      /*tp_weaklistoffset*/
1373         (getiterfunc)BZ2File_getiter, /*tp_iter*/
1374         (iternextfunc)BZ2File_iternext, /*tp_iternext*/
1375         BZ2File_methods,        /*tp_methods*/
1376         0,                      /*tp_members*/
1377         BZ2File_getset,         /*tp_getset*/
1378         0,                      /*tp_base*/
1379         0,                      /*tp_dict*/
1380         0,                      /*tp_descr_get*/
1381         0,                      /*tp_descr_set*/
1382         0,                      /*tp_dictoffset*/
1383         (initproc)BZ2File_init, /*tp_init*/
1384         PyType_GenericAlloc,    /*tp_alloc*/
1385         PyType_GenericNew,      /*tp_new*/
1386         PyObject_Free,          /*tp_free*/
1387         0,                      /*tp_is_gc*/
1388 };
1389
1390
1391 /* ===================================================================== */
1392 /* Methods of BZ2Comp. */
1393
1394 PyDoc_STRVAR(BZ2Comp_compress__doc__,
1395 "compress(data) -> string\n\
1396 \n\
1397 Provide more data to the compressor object. It will return chunks of\n\
1398 compressed data whenever possible. When you've finished providing data\n\
1399 to compress, call the flush() method to finish the compression process,\n\
1400 and return what is left in the internal buffers.\n\
1401 ");
1402
1403 static PyObject *
1404 BZ2Comp_compress(BZ2CompObject *self, PyObject *args)
1405 {
1406         Py_buffer pdata;
1407         char *data;
1408         int datasize;
1409         int bufsize = SMALLCHUNK;
1410         PY_LONG_LONG totalout;
1411         PyObject *ret = NULL;
1412         bz_stream *bzs = &self->bzs;
1413         int bzerror;
1414
1415         if (!PyArg_ParseTuple(args, "y*:compress", &pdata))
1416                 return NULL;
1417         data = pdata.buf;
1418         datasize = pdata.len;
1419
1420         if (datasize == 0) {
1421                 PyBuffer_Release(&pdata);
1422                 return PyBytes_FromStringAndSize("", 0);
1423         }
1424
1425         ACQUIRE_LOCK(self);
1426         if (!self->running) {
1427                 PyErr_SetString(PyExc_ValueError,
1428                                 "this object was already flushed");
1429                 goto error;
1430         }
1431
1432         ret = PyBytes_FromStringAndSize(NULL, bufsize);
1433         if (!ret)
1434                 goto error;
1435
1436         bzs->next_in = data;
1437         bzs->avail_in = datasize;
1438         bzs->next_out = BUF(ret);
1439         bzs->avail_out = bufsize;
1440
1441         totalout = BZS_TOTAL_OUT(bzs);
1442
1443         for (;;) {
1444                 Py_BEGIN_ALLOW_THREADS
1445                 bzerror = BZ2_bzCompress(bzs, BZ_RUN);
1446                 Py_END_ALLOW_THREADS
1447                 if (bzerror != BZ_RUN_OK) {
1448                         Util_CatchBZ2Error(bzerror);
1449                         goto error;
1450                 }
1451                 if (bzs->avail_in == 0)
1452                         break; /* no more input data */
1453                 if (bzs->avail_out == 0) {
1454                         bufsize = Util_NewBufferSize(bufsize);
1455                         if (_PyBytes_Resize(&ret, bufsize) < 0) {
1456                                 BZ2_bzCompressEnd(bzs);
1457                                 goto error;
1458                         }
1459                         bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1460                                                     - totalout);
1461                         bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1462                 }
1463         }
1464
1465         if (_PyBytes_Resize(&ret,
1466                            (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1467                 goto error;
1468
1469         RELEASE_LOCK(self);
1470         PyBuffer_Release(&pdata);
1471         return ret;
1472
1473 error:
1474         RELEASE_LOCK(self);
1475         PyBuffer_Release(&pdata);
1476         Py_XDECREF(ret);
1477         return NULL;
1478 }
1479
1480 PyDoc_STRVAR(BZ2Comp_flush__doc__,
1481 "flush() -> string\n\
1482 \n\
1483 Finish the compression process and return what is left in internal buffers.\n\
1484 You must not use the compressor object after calling this method.\n\
1485 ");
1486
1487 static PyObject *
1488 BZ2Comp_flush(BZ2CompObject *self)
1489 {
1490         int bufsize = SMALLCHUNK;
1491         PyObject *ret = NULL;
1492         bz_stream *bzs = &self->bzs;
1493         PY_LONG_LONG totalout;
1494         int bzerror;
1495
1496         ACQUIRE_LOCK(self);
1497         if (!self->running) {
1498                 PyErr_SetString(PyExc_ValueError, "object was already "
1499                                                   "flushed");
1500                 goto error;
1501         }
1502         self->running = 0;
1503
1504         ret = PyBytes_FromStringAndSize(NULL, bufsize);
1505         if (!ret)
1506                 goto error;
1507
1508         bzs->next_out = BUF(ret);
1509         bzs->avail_out = bufsize;
1510
1511         totalout = BZS_TOTAL_OUT(bzs);
1512
1513         for (;;) {
1514                 Py_BEGIN_ALLOW_THREADS
1515                 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1516                 Py_END_ALLOW_THREADS
1517                 if (bzerror == BZ_STREAM_END) {
1518                         break;
1519                 } else if (bzerror != BZ_FINISH_OK) {
1520                         Util_CatchBZ2Error(bzerror);
1521                         goto error;
1522                 }
1523                 if (bzs->avail_out == 0) {
1524                         bufsize = Util_NewBufferSize(bufsize);
1525                         if (_PyBytes_Resize(&ret, bufsize) < 0)
1526                                 goto error;
1527                         bzs->next_out = BUF(ret);
1528                         bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1529                                                     - totalout);
1530                         bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1531                 }
1532         }
1533
1534         if (bzs->avail_out != 0) {
1535                 if (_PyBytes_Resize(&ret,
1536                             (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1537                         goto error;
1538         }
1539
1540         RELEASE_LOCK(self);
1541         return ret;
1542
1543 error:
1544         RELEASE_LOCK(self);
1545         Py_XDECREF(ret);
1546         return NULL;
1547 }
1548
1549 static PyMethodDef BZ2Comp_methods[] = {
1550         {"compress", (PyCFunction)BZ2Comp_compress, METH_VARARGS,
1551          BZ2Comp_compress__doc__},
1552         {"flush", (PyCFunction)BZ2Comp_flush, METH_NOARGS,
1553          BZ2Comp_flush__doc__},
1554         {NULL,          NULL}           /* sentinel */
1555 };
1556
1557
1558 /* ===================================================================== */
1559 /* Slot definitions for BZ2Comp_Type. */
1560
1561 static int
1562 BZ2Comp_init(BZ2CompObject *self, PyObject *args, PyObject *kwargs)
1563 {
1564         int compresslevel = 9;
1565         int bzerror;
1566         static char *kwlist[] = {"compresslevel", 0};
1567
1568         if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:BZ2Compressor",
1569                                          kwlist, &compresslevel))
1570                 return -1;
1571
1572         if (compresslevel < 1 || compresslevel > 9) {
1573                 PyErr_SetString(PyExc_ValueError,
1574                                 "compresslevel must be between 1 and 9");
1575                 goto error;
1576         }
1577
1578 #ifdef WITH_THREAD
1579         self->lock = PyThread_allocate_lock();
1580         if (!self->lock) {
1581                 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1582                 goto error;
1583         }
1584 #endif
1585
1586         memset(&self->bzs, 0, sizeof(bz_stream));
1587         bzerror = BZ2_bzCompressInit(&self->bzs, compresslevel, 0, 0);
1588         if (bzerror != BZ_OK) {
1589                 Util_CatchBZ2Error(bzerror);
1590                 goto error;
1591         }
1592
1593         self->running = 1;
1594
1595         return 0;
1596 error:
1597 #ifdef WITH_THREAD
1598         if (self->lock) {
1599                 PyThread_free_lock(self->lock);
1600                 self->lock = NULL;
1601         }
1602 #endif
1603         return -1;
1604 }
1605
1606 static void
1607 BZ2Comp_dealloc(BZ2CompObject *self)
1608 {
1609 #ifdef WITH_THREAD
1610         if (self->lock)
1611                 PyThread_free_lock(self->lock);
1612 #endif
1613         BZ2_bzCompressEnd(&self->bzs);
1614         Py_TYPE(self)->tp_free((PyObject *)self);
1615 }
1616
1617
1618 /* ===================================================================== */
1619 /* BZ2Comp_Type definition. */
1620
1621 PyDoc_STRVAR(BZ2Comp__doc__,
1622 "BZ2Compressor([compresslevel=9]) -> compressor object\n\
1623 \n\
1624 Create a new compressor object. This object may be used to compress\n\
1625 data sequentially. If you want to compress data in one shot, use the\n\
1626 compress() function instead. The compresslevel parameter, if given,\n\
1627 must be a number between 1 and 9.\n\
1628 ");
1629
1630 static PyTypeObject BZ2Comp_Type = {
1631         PyVarObject_HEAD_INIT(NULL, 0)
1632         "bz2.BZ2Compressor",    /*tp_name*/
1633         sizeof(BZ2CompObject),  /*tp_basicsize*/
1634         0,                      /*tp_itemsize*/
1635         (destructor)BZ2Comp_dealloc, /*tp_dealloc*/
1636         0,                      /*tp_print*/
1637         0,                      /*tp_getattr*/
1638         0,                      /*tp_setattr*/
1639         0,                      /*tp_reserved*/
1640         0,                      /*tp_repr*/
1641         0,                      /*tp_as_number*/
1642         0,                      /*tp_as_sequence*/
1643         0,                      /*tp_as_mapping*/
1644         0,                      /*tp_hash*/
1645         0,                      /*tp_call*/
1646         0,                      /*tp_str*/
1647         PyObject_GenericGetAttr,/*tp_getattro*/
1648         PyObject_GenericSetAttr,/*tp_setattro*/
1649         0,                      /*tp_as_buffer*/
1650         Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1651         BZ2Comp__doc__,         /*tp_doc*/
1652         0,                      /*tp_traverse*/
1653         0,                      /*tp_clear*/
1654         0,                      /*tp_richcompare*/
1655         0,                      /*tp_weaklistoffset*/
1656         0,                      /*tp_iter*/
1657         0,                      /*tp_iternext*/
1658         BZ2Comp_methods,        /*tp_methods*/
1659         0,                      /*tp_members*/
1660         0,                      /*tp_getset*/
1661         0,                      /*tp_base*/
1662         0,                      /*tp_dict*/
1663         0,                      /*tp_descr_get*/
1664         0,                      /*tp_descr_set*/
1665         0,                      /*tp_dictoffset*/
1666         (initproc)BZ2Comp_init, /*tp_init*/
1667         PyType_GenericAlloc,    /*tp_alloc*/
1668         PyType_GenericNew,      /*tp_new*/
1669         PyObject_Free,          /*tp_free*/
1670         0,                      /*tp_is_gc*/
1671 };
1672
1673
1674 /* ===================================================================== */
1675 /* Members of BZ2Decomp. */
1676
1677 #undef OFF
1678 #define OFF(x) offsetof(BZ2DecompObject, x)
1679
1680 static PyMemberDef BZ2Decomp_members[] = {
1681         {"unused_data", T_OBJECT, OFF(unused_data), READONLY},
1682         {NULL}  /* Sentinel */
1683 };
1684
1685
1686 /* ===================================================================== */
1687 /* Methods of BZ2Decomp. */
1688
1689 PyDoc_STRVAR(BZ2Decomp_decompress__doc__,
1690 "decompress(data) -> string\n\
1691 \n\
1692 Provide more data to the decompressor object. It will return chunks\n\
1693 of decompressed data whenever possible. If you try to decompress data\n\
1694 after the end of stream is found, EOFError will be raised. If any data\n\
1695 was found after the end of stream, it'll be ignored and saved in\n\
1696 unused_data attribute.\n\
1697 ");
1698
1699 static PyObject *
1700 BZ2Decomp_decompress(BZ2DecompObject *self, PyObject *args)
1701 {
1702         Py_buffer pdata;
1703         char *data;
1704         int datasize;
1705         int bufsize = SMALLCHUNK;
1706         PY_LONG_LONG totalout;
1707         PyObject *ret = NULL;
1708         bz_stream *bzs = &self->bzs;
1709         int bzerror;
1710
1711         if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
1712                 return NULL;
1713         data = pdata.buf;
1714         datasize = pdata.len;
1715
1716         ACQUIRE_LOCK(self);
1717         if (!self->running) {
1718                 PyErr_SetString(PyExc_EOFError, "end of stream was "
1719                                                 "already found");
1720                 goto error;
1721         }
1722
1723         ret = PyBytes_FromStringAndSize(NULL, bufsize);
1724         if (!ret)
1725                 goto error;
1726
1727         bzs->next_in = data;
1728         bzs->avail_in = datasize;
1729         bzs->next_out = BUF(ret);
1730         bzs->avail_out = bufsize;
1731
1732         totalout = BZS_TOTAL_OUT(bzs);
1733
1734         for (;;) {
1735                 Py_BEGIN_ALLOW_THREADS
1736                 bzerror = BZ2_bzDecompress(bzs);
1737                 Py_END_ALLOW_THREADS
1738                 if (bzerror == BZ_STREAM_END) {
1739                         if (bzs->avail_in != 0) {
1740                                 Py_DECREF(self->unused_data);
1741                                 self->unused_data =
1742                                     PyBytes_FromStringAndSize(bzs->next_in,
1743                                                                bzs->avail_in);
1744                         }
1745                         self->running = 0;
1746                         break;
1747                 }
1748                 if (bzerror != BZ_OK) {
1749                         Util_CatchBZ2Error(bzerror);
1750                         goto error;
1751                 }
1752                 if (bzs->avail_in == 0)
1753                         break; /* no more input data */
1754                 if (bzs->avail_out == 0) {
1755                         bufsize = Util_NewBufferSize(bufsize);
1756                         if (_PyBytes_Resize(&ret, bufsize) < 0) {
1757                                 BZ2_bzDecompressEnd(bzs);
1758                                 goto error;
1759                         }
1760                         bzs->next_out = BUF(ret);
1761                         bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1762                                                     - totalout);
1763                         bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1764                 }
1765         }
1766
1767         if (bzs->avail_out != 0) {
1768                 if (_PyBytes_Resize(&ret,
1769                             (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1770                         goto error;
1771         }
1772
1773         RELEASE_LOCK(self);
1774         PyBuffer_Release(&pdata);
1775         return ret;
1776
1777 error:
1778         RELEASE_LOCK(self);
1779         PyBuffer_Release(&pdata);
1780         Py_XDECREF(ret);
1781         return NULL;
1782 }
1783
1784 static PyMethodDef BZ2Decomp_methods[] = {
1785         {"decompress", (PyCFunction)BZ2Decomp_decompress, METH_VARARGS, BZ2Decomp_decompress__doc__},
1786         {NULL,          NULL}           /* sentinel */
1787 };
1788
1789
1790 /* ===================================================================== */
1791 /* Slot definitions for BZ2Decomp_Type. */
1792
1793 static int
1794 BZ2Decomp_init(BZ2DecompObject *self, PyObject *args, PyObject *kwargs)
1795 {
1796         int bzerror;
1797
1798         if (!PyArg_ParseTuple(args, ":BZ2Decompressor"))
1799                 return -1;
1800
1801 #ifdef WITH_THREAD
1802         self->lock = PyThread_allocate_lock();
1803         if (!self->lock) {
1804                 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1805                 goto error;
1806         }
1807 #endif
1808
1809         self->unused_data = PyBytes_FromStringAndSize("", 0);
1810         if (!self->unused_data)
1811                 goto error;
1812
1813         memset(&self->bzs, 0, sizeof(bz_stream));
1814         bzerror = BZ2_bzDecompressInit(&self->bzs, 0, 0);
1815         if (bzerror != BZ_OK) {
1816                 Util_CatchBZ2Error(bzerror);
1817                 goto error;
1818         }
1819
1820         self->running = 1;
1821
1822         return 0;
1823
1824 error:
1825 #ifdef WITH_THREAD
1826         if (self->lock) {
1827                 PyThread_free_lock(self->lock);
1828                 self->lock = NULL;
1829         }
1830 #endif
1831         Py_CLEAR(self->unused_data);
1832         return -1;
1833 }
1834
1835 static void
1836 BZ2Decomp_dealloc(BZ2DecompObject *self)
1837 {
1838 #ifdef WITH_THREAD
1839         if (self->lock)
1840                 PyThread_free_lock(self->lock);
1841 #endif
1842         Py_XDECREF(self->unused_data);
1843         BZ2_bzDecompressEnd(&self->bzs);
1844         Py_TYPE(self)->tp_free((PyObject *)self);
1845 }
1846
1847
1848 /* ===================================================================== */
1849 /* BZ2Decomp_Type definition. */
1850
1851 PyDoc_STRVAR(BZ2Decomp__doc__,
1852 "BZ2Decompressor() -> decompressor object\n\
1853 \n\
1854 Create a new decompressor object. This object may be used to decompress\n\
1855 data sequentially. If you want to decompress data in one shot, use the\n\
1856 decompress() function instead.\n\
1857 ");
1858
1859 static PyTypeObject BZ2Decomp_Type = {
1860         PyVarObject_HEAD_INIT(NULL, 0)
1861         "bz2.BZ2Decompressor",  /*tp_name*/
1862         sizeof(BZ2DecompObject), /*tp_basicsize*/
1863         0,                      /*tp_itemsize*/
1864         (destructor)BZ2Decomp_dealloc, /*tp_dealloc*/
1865         0,                      /*tp_print*/
1866         0,                      /*tp_getattr*/
1867         0,                      /*tp_setattr*/
1868         0,                      /*tp_reserved*/
1869         0,                      /*tp_repr*/
1870         0,                      /*tp_as_number*/
1871         0,                      /*tp_as_sequence*/
1872         0,                      /*tp_as_mapping*/
1873         0,                      /*tp_hash*/
1874         0,                      /*tp_call*/
1875         0,                      /*tp_str*/
1876         PyObject_GenericGetAttr,/*tp_getattro*/
1877         PyObject_GenericSetAttr,/*tp_setattro*/
1878         0,                      /*tp_as_buffer*/
1879         Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1880         BZ2Decomp__doc__,       /*tp_doc*/
1881         0,                      /*tp_traverse*/
1882         0,                      /*tp_clear*/
1883         0,                      /*tp_richcompare*/
1884         0,                      /*tp_weaklistoffset*/
1885         0,                      /*tp_iter*/
1886         0,                      /*tp_iternext*/
1887         BZ2Decomp_methods,      /*tp_methods*/
1888         BZ2Decomp_members,      /*tp_members*/
1889         0,                      /*tp_getset*/
1890         0,                      /*tp_base*/
1891         0,                      /*tp_dict*/
1892         0,                      /*tp_descr_get*/
1893         0,                      /*tp_descr_set*/
1894         0,                      /*tp_dictoffset*/
1895         (initproc)BZ2Decomp_init, /*tp_init*/
1896         PyType_GenericAlloc,    /*tp_alloc*/
1897         PyType_GenericNew,      /*tp_new*/
1898         PyObject_Free,          /*tp_free*/
1899         0,                      /*tp_is_gc*/
1900 };
1901
1902
1903 /* ===================================================================== */
1904 /* Module functions. */
1905
1906 PyDoc_STRVAR(bz2_compress__doc__,
1907 "compress(data [, compresslevel=9]) -> string\n\
1908 \n\
1909 Compress data in one shot. If you want to compress data sequentially,\n\
1910 use an instance of BZ2Compressor instead. The compresslevel parameter, if\n\
1911 given, must be a number between 1 and 9.\n\
1912 ");
1913
1914 static PyObject *
1915 bz2_compress(PyObject *self, PyObject *args, PyObject *kwargs)
1916 {
1917         int compresslevel=9;
1918         Py_buffer pdata;
1919         char *data;
1920         int datasize;
1921         int bufsize;
1922         PyObject *ret = NULL;
1923         bz_stream _bzs;
1924         bz_stream *bzs = &_bzs;
1925         int bzerror;
1926         static char *kwlist[] = {"data", "compresslevel", 0};
1927
1928         if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|i",
1929                                          kwlist, &pdata,
1930                                          &compresslevel))
1931                 return NULL;
1932         data = pdata.buf;
1933         datasize = pdata.len;
1934
1935         if (compresslevel < 1 || compresslevel > 9) {
1936                 PyErr_SetString(PyExc_ValueError,
1937                                 "compresslevel must be between 1 and 9");
1938                 PyBuffer_Release(&pdata);
1939                 return NULL;
1940         }
1941
1942         /* Conforming to bz2 manual, this is large enough to fit compressed
1943          * data in one shot. We will check it later anyway. */
1944         bufsize = datasize + (datasize/100+1) + 600;
1945
1946         ret = PyBytes_FromStringAndSize(NULL, bufsize);
1947         if (!ret) {
1948                 PyBuffer_Release(&pdata);
1949                 return NULL;
1950         }
1951
1952         memset(bzs, 0, sizeof(bz_stream));
1953
1954         bzs->next_in = data;
1955         bzs->avail_in = datasize;
1956         bzs->next_out = BUF(ret);
1957         bzs->avail_out = bufsize;
1958
1959         bzerror = BZ2_bzCompressInit(bzs, compresslevel, 0, 0);
1960         if (bzerror != BZ_OK) {
1961                 Util_CatchBZ2Error(bzerror);
1962                 PyBuffer_Release(&pdata);
1963                 Py_DECREF(ret);
1964                 return NULL;
1965         }
1966
1967         for (;;) {
1968                 Py_BEGIN_ALLOW_THREADS
1969                 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1970                 Py_END_ALLOW_THREADS
1971                 if (bzerror == BZ_STREAM_END) {
1972                         break;
1973                 } else if (bzerror != BZ_FINISH_OK) {
1974                         BZ2_bzCompressEnd(bzs);
1975                         Util_CatchBZ2Error(bzerror);
1976                         PyBuffer_Release(&pdata);
1977                         Py_DECREF(ret);
1978                         return NULL;
1979                 }
1980                 if (bzs->avail_out == 0) {
1981                         bufsize = Util_NewBufferSize(bufsize);
1982                         if (_PyBytes_Resize(&ret, bufsize) < 0) {
1983                                 BZ2_bzCompressEnd(bzs);
1984                                 PyBuffer_Release(&pdata);
1985                                 return NULL;
1986                         }
1987                         bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
1988                         bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1989                 }
1990         }
1991
1992         if (bzs->avail_out != 0) {
1993                 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
1994                         ret = NULL;
1995                 }
1996         }
1997         BZ2_bzCompressEnd(bzs);
1998
1999         PyBuffer_Release(&pdata);
2000         return ret;
2001 }
2002
2003 PyDoc_STRVAR(bz2_decompress__doc__,
2004 "decompress(data) -> decompressed data\n\
2005 \n\
2006 Decompress data in one shot. If you want to decompress data sequentially,\n\
2007 use an instance of BZ2Decompressor instead.\n\
2008 ");
2009
2010 static PyObject *
2011 bz2_decompress(PyObject *self, PyObject *args)
2012 {
2013         Py_buffer pdata;
2014         char *data;
2015         int datasize;
2016         int bufsize = SMALLCHUNK;
2017         PyObject *ret;
2018         bz_stream _bzs;
2019         bz_stream *bzs = &_bzs;
2020         int bzerror;
2021
2022         if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
2023                 return NULL;
2024         data = pdata.buf;
2025         datasize = pdata.len;
2026
2027         if (datasize == 0) {
2028                 PyBuffer_Release(&pdata);
2029                 return PyBytes_FromStringAndSize("", 0);
2030         }
2031
2032         ret = PyBytes_FromStringAndSize(NULL, bufsize);
2033         if (!ret) {
2034                 PyBuffer_Release(&pdata);
2035                 return NULL;
2036         }
2037
2038         memset(bzs, 0, sizeof(bz_stream));
2039
2040         bzs->next_in = data;
2041         bzs->avail_in = datasize;
2042         bzs->next_out = BUF(ret);
2043         bzs->avail_out = bufsize;
2044
2045         bzerror = BZ2_bzDecompressInit(bzs, 0, 0);
2046         if (bzerror != BZ_OK) {
2047                 Util_CatchBZ2Error(bzerror);
2048                 Py_DECREF(ret);
2049                 PyBuffer_Release(&pdata);
2050                 return NULL;
2051         }
2052
2053         for (;;) {
2054                 Py_BEGIN_ALLOW_THREADS
2055                 bzerror = BZ2_bzDecompress(bzs);
2056                 Py_END_ALLOW_THREADS
2057                 if (bzerror == BZ_STREAM_END) {
2058                         break;
2059                 } else if (bzerror != BZ_OK) {
2060                         BZ2_bzDecompressEnd(bzs);
2061                         Util_CatchBZ2Error(bzerror);
2062                         PyBuffer_Release(&pdata);
2063                         Py_DECREF(ret);
2064                         return NULL;
2065                 }
2066                 if (bzs->avail_in == 0) {
2067                         BZ2_bzDecompressEnd(bzs);
2068                         PyErr_SetString(PyExc_ValueError,
2069                                         "couldn't find end of stream");
2070                         PyBuffer_Release(&pdata);
2071                         Py_DECREF(ret);
2072                         return NULL;
2073                 }
2074                 if (bzs->avail_out == 0) {
2075                         bufsize = Util_NewBufferSize(bufsize);
2076                         if (_PyBytes_Resize(&ret, bufsize) < 0) {
2077                                 BZ2_bzDecompressEnd(bzs);
2078                                 PyBuffer_Release(&pdata);
2079                                 return NULL;
2080                         }
2081                         bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
2082                         bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
2083                 }
2084         }
2085
2086         if (bzs->avail_out != 0) {
2087                 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
2088                         ret = NULL;
2089                 }
2090         }
2091         BZ2_bzDecompressEnd(bzs);
2092         PyBuffer_Release(&pdata);
2093
2094         return ret;
2095 }
2096
2097 static PyMethodDef bz2_methods[] = {
2098         {"compress", (PyCFunction) bz2_compress, METH_VARARGS|METH_KEYWORDS,
2099                 bz2_compress__doc__},
2100         {"decompress", (PyCFunction) bz2_decompress, METH_VARARGS,
2101                 bz2_decompress__doc__},
2102         {NULL,          NULL}           /* sentinel */
2103 };
2104
2105 /* ===================================================================== */
2106 /* Initialization function. */
2107
2108 PyDoc_STRVAR(bz2__doc__,
2109 "The python bz2 module provides a comprehensive interface for\n\
2110 the bz2 compression library. It implements a complete file\n\
2111 interface, one shot (de)compression functions, and types for\n\
2112 sequential (de)compression.\n\
2113 ");
2114
2115
2116 static struct PyModuleDef bz2module = {
2117         PyModuleDef_HEAD_INIT,
2118         "bz2",
2119         bz2__doc__,
2120         -1,
2121         bz2_methods,
2122         NULL,
2123         NULL,
2124         NULL,
2125         NULL
2126 };
2127
2128 PyMODINIT_FUNC
2129 PyInit_bz2(void)
2130 {
2131         PyObject *m;
2132
2133         Py_TYPE(&BZ2File_Type) = &PyType_Type;
2134         Py_TYPE(&BZ2Comp_Type) = &PyType_Type;
2135         Py_TYPE(&BZ2Decomp_Type) = &PyType_Type;
2136
2137         m = PyModule_Create(&bz2module);
2138         if (m == NULL)
2139                 return NULL;
2140
2141         PyModule_AddObject(m, "__author__", PyUnicode_FromString(__author__));
2142
2143         Py_INCREF(&BZ2File_Type);
2144         PyModule_AddObject(m, "BZ2File", (PyObject *)&BZ2File_Type);
2145
2146         Py_INCREF(&BZ2Comp_Type);
2147         PyModule_AddObject(m, "BZ2Compressor", (PyObject *)&BZ2Comp_Type);
2148
2149         Py_INCREF(&BZ2Decomp_Type);
2150         PyModule_AddObject(m, "BZ2Decompressor", (PyObject *)&BZ2Decomp_Type);
2151         return m;
2152 }