Modules/zlibmodule.c

   1 /* zlibmodule.c -- gzip-compatible data compression */
   2 /* See http://www.gzip.org/zlib/ */
   3
   4 /* Windows users:  read Python's PCbuild\readme.txt */
   5
   6
   7 #include "Python.h"
   8 #include "zlib.h"
   9
  10 #ifdef WITH_THREAD
  11 #include "pythread.h"
  12
  13 /* #defs ripped off from _tkinter.c, even though the situation here is much
  14    simpler, because we don't have to worry about waiting for Tcl
  15    events!  And, since zlib itself is threadsafe, we don't need to worry
  16    about re-entering zlib functions.
  17
  18    N.B.
  19
  20    Since ENTER_ZLIB and LEAVE_ZLIB only need to be called on functions
  21    that modify the components of preexisting de/compress objects, it
  22    could prove to be a performance gain on multiprocessor machines if
  23    there was an de/compress object-specific lock.  However, for the
  24    moment the ENTER_ZLIB and LEAVE_ZLIB calls are global for ALL
  25    de/compress objects.
  26  */
  27
  28 static PyThread_type_lock zlib_lock = NULL; /* initialized on module load */
  29
  30 #define ENTER_ZLIB \
  31         Py_BEGIN_ALLOW_THREADS \
  32         PyThread_acquire_lock(zlib_lock, 1); \
  33         Py_END_ALLOW_THREADS
  34
  35 #define LEAVE_ZLIB \
  36         PyThread_release_lock(zlib_lock);
  37
  38 #else
  39
  40 #define ENTER_ZLIB
  41 #define LEAVE_ZLIB
  42
  43 #endif
  44
  45 /* The following parameters are copied from zutil.h, version 0.95 */
  46 #define DEFLATED   8
  47 #if MAX_MEM_LEVEL >= 8
  48 #  define DEF_MEM_LEVEL 8
  49 #else
  50 #  define DEF_MEM_LEVEL  MAX_MEM_LEVEL
  51 #endif
  52 #define DEF_WBITS MAX_WBITS
  53
  54 /* The output buffer will be increased in chunks of DEFAULTALLOC bytes. */
  55 #define DEFAULTALLOC (16*1024)
  56 #define PyInit_zlib initzlib
  57
  58 static PyTypeObject Comptype;
  59 static PyTypeObject Decomptype;
  60
  61 static PyObject *ZlibError;
  62
  63 typedef struct
  64 {
  65     PyObject_HEAD
  66     z_stream zst;
  67     PyObject *unused_data;
  68     PyObject *unconsumed_tail;
  69     int is_initialised;
  70 } compobject;
  71
  72 static void
  73 zlib_error(z_stream zst, int err, char *msg)
  74 {
  75     if (zst.msg == Z_NULL)
  76         PyErr_Format(ZlibError, "Error %d %s", err, msg);
  77     else
  78         PyErr_Format(ZlibError, "Error %d %s: %.200s", err, msg, zst.msg);
  79 }
  80
  81 PyDoc_STRVAR(compressobj__doc__,
  82 "compressobj([level]) -- Return a compressor object.\n"
  83 "\n"
  84 "Optional arg level is the compression level, in 1-9.");
  85
  86 PyDoc_STRVAR(decompressobj__doc__,
  87 "decompressobj([wbits]) -- Return a decompressor object.\n"
  88 "\n"
  89 "Optional arg wbits is the window buffer size.");
  90
  91 static compobject *
  92 newcompobject(PyTypeObject *type)
  93 {
  94     compobject *self;
  95     self = PyObject_New(compobject, type);
  96     if (self == NULL)
  97         return NULL;
  98     self->is_initialised = 0;
  99     self->unused_data = PyString_FromString("");
 100     if (self->unused_data == NULL) {
 101         Py_DECREF(self);
 102         return NULL;
 103     }
 104     self->unconsumed_tail = PyString_FromString("");
 105     if (self->unconsumed_tail == NULL) {
 106         Py_DECREF(self);
 107         return NULL;
 108     }
 109     return self;
 110 }
 111
 112 PyDoc_STRVAR(compress__doc__,
 113 "compress(string[, level]) -- Returned compressed string.\n"
 114 "\n"
 115 "Optional arg level is the compression level, in 1-9.");
 116
 117 static PyObject *
 118 PyZlib_compress(PyObject *self, PyObject *args)
 119 {
 120     PyObject *ReturnVal = NULL;
 121     Byte *input, *output;
 122     int length, level=Z_DEFAULT_COMPRESSION, err;
 123     z_stream zst;
 124
 125     /* require Python string object, optional 'level' arg */
 126     if (!PyArg_ParseTuple(args, "s#|i:compress", &input, &length, &level))
 127         return NULL;
 128
 129     zst.avail_out = length + length/1000 + 12 + 1;
 130
 131     output = (Byte*)malloc(zst.avail_out);
 132     if (output == NULL) {
 133         PyErr_SetString(PyExc_MemoryError,
 134                         "Can't allocate memory to compress data");
 135         return NULL;
 136     }
 137
 138     /* Past the point of no return.  From here on out, we need to make sure
 139        we clean up mallocs & INCREFs. */
 140
 141     zst.zalloc = (alloc_func)NULL;
 142     zst.zfree = (free_func)Z_NULL;
 143     zst.next_out = (Byte *)output;
 144     zst.next_in = (Byte *)input;
 145     zst.avail_in = length;
 146     err = deflateInit(&zst, level);
 147
 148     switch(err) {
 149     case(Z_OK):
 150         break;
 151     case(Z_MEM_ERROR):
 152         PyErr_SetString(PyExc_MemoryError,
 153                         "Out of memory while compressing data");
 154         goto error;
 155     case(Z_STREAM_ERROR):
 156         PyErr_SetString(ZlibError,
 157                         "Bad compression level");
 158         goto error;
 159     default:
 160         deflateEnd(&zst);
 161         zlib_error(zst, err, "while compressing data");
 162         goto error;
 163     }
 164
 165     Py_BEGIN_ALLOW_THREADS;
 166     err = deflate(&zst, Z_FINISH);
 167     Py_END_ALLOW_THREADS;
 168
 169     if (err != Z_STREAM_END) {
 170         zlib_error(zst, err, "while compressing data");
 171         deflateEnd(&zst);
 172         goto error;
 173     }
 174
 175     err=deflateEnd(&zst);
 176     if (err == Z_OK)
 177         ReturnVal = PyString_FromStringAndSize((char *)output,
 178                                                zst.total_out);
 179     else
 180         zlib_error(zst, err, "while finishing compression");
 181
 182  error:
 183     free(output);
 184
 185     return ReturnVal;
 186 }
 187
 188 PyDoc_STRVAR(decompress__doc__,
 189 "decompress(string[, wbits[, bufsize]]) -- Return decompressed string.\n"
 190 "\n"
 191 "Optional arg wbits is the window buffer size.  Optional arg bufsize is\n"
 192 "the initial output buffer size.");
 193
 194 static PyObject *
 195 PyZlib_decompress(PyObject *self, PyObject *args)
 196 {
 197     PyObject *result_str;
 198     Byte *input;
 199     int length, err;
 200     int wsize=DEF_WBITS;
 201     Py_ssize_t r_strlen=DEFAULTALLOC;
 202     z_stream zst;
 203
 204     if (!PyArg_ParseTuple(args, "s#|in:decompress",
 205                           &input, &length, &wsize, &r_strlen))
 206         return NULL;
 207
 208     if (r_strlen <= 0)
 209         r_strlen = 1;
 210
 211     zst.avail_in = length;
 212     zst.avail_out = r_strlen;
 213
 214     if (!(result_str = PyString_FromStringAndSize(NULL, r_strlen)))
 215         return NULL;
 216
 217     zst.zalloc = (alloc_func)NULL;
 218     zst.zfree = (free_func)Z_NULL;
 219     zst.next_out = (Byte *)PyString_AS_STRING(result_str);
 220     zst.next_in = (Byte *)input;
 221     err = inflateInit2(&zst, wsize);
 222
 223     switch(err) {
 224     case(Z_OK):
 225         break;
 226     case(Z_MEM_ERROR):
 227         PyErr_SetString(PyExc_MemoryError,
 228                         "Out of memory while decompressing data");
 229         goto error;
 230     default:
 231         inflateEnd(&zst);
 232         zlib_error(zst, err, "while preparing to decompress data");
 233         goto error;
 234     }
 235
 236     do {
 237         Py_BEGIN_ALLOW_THREADS
 238         err=inflate(&zst, Z_FINISH);
 239         Py_END_ALLOW_THREADS
 240
 241         switch(err) {
 242         case(Z_STREAM_END):
 243             break;
 244         case(Z_BUF_ERROR):
 245             /*
 246              * If there is at least 1 byte of room according to zst.avail_out
 247              * and we get this error, assume that it means zlib cannot
 248              * process the inflate call() due to an error in the data.
 249              */
 250             if (zst.avail_out > 0) {
 251                 PyErr_Format(ZlibError, "Error %i while decompressing data",
 252                              err);
 253                 inflateEnd(&zst);
 254                 goto error;
 255             }
 256             /* fall through */
 257         case(Z_OK):
 258             /* need more memory */
 259             if (_PyString_Resize(&result_str, r_strlen << 1) < 0) {
 260                 inflateEnd(&zst);
 261                 goto error;
 262             }
 263             zst.next_out = (unsigned char *)PyString_AS_STRING(result_str) \
 264                 + r_strlen;
 265             zst.avail_out = r_strlen;
 266             r_strlen = r_strlen << 1;
 267             break;
 268         default:
 269             inflateEnd(&zst);
 270             zlib_error(zst, err, "while decompressing data");
 271             goto error;
 272         }
 273     } while (err != Z_STREAM_END);
 274
 275     err = inflateEnd(&zst);
 276     if (err != Z_OK) {
 277         zlib_error(zst, err, "while finishing data decompression");
 278         goto error;
 279     }
 280
 281     _PyString_Resize(&result_str, zst.total_out);
 282     return result_str;
 283
 284  error:
 285     Py_XDECREF(result_str);
 286     return NULL;
 287 }
 288
 289 static PyObject *
 290 PyZlib_compressobj(PyObject *selfptr, PyObject *args)
 291 {
 292     compobject *self;
 293     int level=Z_DEFAULT_COMPRESSION, method=DEFLATED;
 294     int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err;
 295
 296     if (!PyArg_ParseTuple(args, "|iiiii:compressobj", &level, &method, &wbits,
 297                           &memLevel, &strategy))
 298         return NULL;
 299
 300     self = newcompobject(&Comptype);
 301     if (self==NULL)
 302         return(NULL);
 303     self->zst.zalloc = (alloc_func)NULL;
 304     self->zst.zfree = (free_func)Z_NULL;
 305     self->zst.next_in = NULL;
 306     self->zst.avail_in = 0;
 307     err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
 308     switch(err) {
 309     case (Z_OK):
 310         self->is_initialised = 1;
 311         return (PyObject*)self;
 312     case (Z_MEM_ERROR):
 313         Py_DECREF(self);
 314         PyErr_SetString(PyExc_MemoryError,
 315                         "Can't allocate memory for compression object");
 316         return NULL;
 317     case(Z_STREAM_ERROR):
 318         Py_DECREF(self);
 319         PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
 320         return NULL;
 321     default:
 322         zlib_error(self->zst, err, "while creating compression object");
 323         Py_DECREF(self);
 324         return NULL;
 325     }
 326 }
 327
 328 static PyObject *
 329 PyZlib_decompressobj(PyObject *selfptr, PyObject *args)
 330 {
 331     int wbits=DEF_WBITS, err;
 332     compobject *self;
 333     if (!PyArg_ParseTuple(args, "|i:decompressobj", &wbits))
 334         return NULL;
 335
 336     self = newcompobject(&Decomptype);
 337     if (self == NULL)
 338         return(NULL);
 339     self->zst.zalloc = (alloc_func)NULL;
 340     self->zst.zfree = (free_func)Z_NULL;
 341     self->zst.next_in = NULL;
 342     self->zst.avail_in = 0;
 343     err = inflateInit2(&self->zst, wbits);
 344     switch(err) {
 345     case (Z_OK):
 346         self->is_initialised = 1;
 347         return (PyObject*)self;
 348     case(Z_STREAM_ERROR):
 349         Py_DECREF(self);
 350         PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
 351         return NULL;
 352     case (Z_MEM_ERROR):
 353         Py_DECREF(self);
 354         PyErr_SetString(PyExc_MemoryError,
 355                         "Can't allocate memory for decompression object");
 356         return NULL;
 357     default:
 358         zlib_error(self->zst, err, "while creating decompression object");
 359         Py_DECREF(self);
 360         return NULL;
 361     }
 362 }
 363
 364 static void
 365 Comp_dealloc(compobject *self)
 366 {
 367     if (self->is_initialised)
 368         deflateEnd(&self->zst);
 369     Py_XDECREF(self->unused_data);
 370     Py_XDECREF(self->unconsumed_tail);
 371     PyObject_Del(self);
 372 }
 373
 374 static void
 375 Decomp_dealloc(compobject *self)
 376 {
 377     if (self->is_initialised)
 378         inflateEnd(&self->zst);
 379     Py_XDECREF(self->unused_data);
 380     Py_XDECREF(self->unconsumed_tail);
 381     PyObject_Del(self);
 382 }
 383
 384 PyDoc_STRVAR(comp_compress__doc__,
 385 "compress(data) -- Return a string containing data compressed.\n"
 386 "\n"
 387 "After calling this function, some of the input data may still\n"
 388 "be stored in internal buffers for later processing.\n"
 389 "Call the flush() method to clear these buffers.");
 390
 391
 392 static PyObject *
 393 PyZlib_objcompress(compobject *self, PyObject *args)
 394 {
 395     int err, inplen, length = DEFAULTALLOC;
 396     PyObject *RetVal;
 397     Byte *input;
 398     unsigned long start_total_out;
 399
 400     if (!PyArg_ParseTuple(args, "s#:compress", &input, &inplen))
 401         return NULL;
 402
 403     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 404         return NULL;
 405
 406     ENTER_ZLIB
 407
 408     start_total_out = self->zst.total_out;
 409     self->zst.avail_in = inplen;
 410     self->zst.next_in = input;
 411     self->zst.avail_out = length;
 412     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 413
 414     Py_BEGIN_ALLOW_THREADS
 415     err = deflate(&(self->zst), Z_NO_FLUSH);
 416     Py_END_ALLOW_THREADS
 417
 418     /* while Z_OK and the output buffer is full, there might be more output,
 419        so extend the output buffer and try again */
 420     while (err == Z_OK && self->zst.avail_out == 0) {
 421         if (_PyString_Resize(&RetVal, length << 1) < 0)
 422             goto error;
 423         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 424             + length;
 425         self->zst.avail_out = length;
 426         length = length << 1;
 427
 428         Py_BEGIN_ALLOW_THREADS
 429         err = deflate(&(self->zst), Z_NO_FLUSH);
 430         Py_END_ALLOW_THREADS
 431     }
 432     /* We will only get Z_BUF_ERROR if the output buffer was full but
 433        there wasn't more output when we tried again, so it is not an error
 434        condition.
 435     */
 436
 437     if (err != Z_OK && err != Z_BUF_ERROR) {
 438         zlib_error(self->zst, err, "while compressing");
 439         Py_DECREF(RetVal);
 440         RetVal = NULL;
 441         goto error;
 442     }
 443     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 444
 445  error:
 446     LEAVE_ZLIB
 447     return RetVal;
 448 }
 449
 450 PyDoc_STRVAR(decomp_decompress__doc__,
 451 "decompress(data, max_length) -- Return a string containing the decompressed\n"
 452 "version of the data.\n"
 453 "\n"
 454 "After calling this function, some of the input data may still be stored in\n"
 455 "internal buffers for later processing.\n"
 456 "Call the flush() method to clear these buffers.\n"
 457 "If the max_length parameter is specified then the return value will be\n"
 458 "no longer than max_length.  Unconsumed input data will be stored in\n"
 459 "the unconsumed_tail attribute.");
 460
 461 static PyObject *
 462 PyZlib_objdecompress(compobject *self, PyObject *args)
 463 {
 464     int err, inplen, old_length, length = DEFAULTALLOC;
 465     int max_length = 0;
 466     PyObject *RetVal;
 467     Byte *input;
 468     unsigned long start_total_out;
 469
 470     if (!PyArg_ParseTuple(args, "s#|i:decompress", &input,
 471                           &inplen, &max_length))
 472         return NULL;
 473     if (max_length < 0) {
 474         PyErr_SetString(PyExc_ValueError,
 475                         "max_length must be greater than zero");
 476         return NULL;
 477     }
 478
 479     /* limit amount of data allocated to max_length */
 480     if (max_length && length > max_length)
 481         length = max_length;
 482     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 483         return NULL;
 484
 485     ENTER_ZLIB
 486
 487     start_total_out = self->zst.total_out;
 488     self->zst.avail_in = inplen;
 489     self->zst.next_in = input;
 490     self->zst.avail_out = length;
 491     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 492
 493     Py_BEGIN_ALLOW_THREADS
 494     err = inflate(&(self->zst), Z_SYNC_FLUSH);
 495     Py_END_ALLOW_THREADS
 496
 497     /* While Z_OK and the output buffer is full, there might be more output.
 498        So extend the output buffer and try again.
 499     */
 500     while (err == Z_OK && self->zst.avail_out == 0) {
 501         /* If max_length set, don't continue decompressing if we've already
 502            reached the limit.
 503         */
 504         if (max_length && length >= max_length)
 505             break;
 506
 507         /* otherwise, ... */
 508         old_length = length;
 509         length = length << 1;
 510         if (max_length && length > max_length)
 511             length = max_length;
 512
 513         if (_PyString_Resize(&RetVal, length) < 0)
 514             goto error;
 515         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 516             + old_length;
 517         self->zst.avail_out = length - old_length;
 518
 519         Py_BEGIN_ALLOW_THREADS
 520         err = inflate(&(self->zst), Z_SYNC_FLUSH);
 521         Py_END_ALLOW_THREADS
 522     }
 523
 524     /* Not all of the compressed data could be accommodated in the output buffer
 525        of specified size. Return the unconsumed tail in an attribute.*/
 526     if(max_length) {
 527         Py_DECREF(self->unconsumed_tail);
 528         self->unconsumed_tail = PyString_FromStringAndSize((char *)self->zst.next_in,
 529                                                            self->zst.avail_in);
 530         if(!self->unconsumed_tail) {
 531             Py_DECREF(RetVal);
 532             RetVal = NULL;
 533             goto error;
 534         }
 535     }
 536
 537     /* The end of the compressed data has been reached, so set the
 538        unused_data attribute to a string containing the remainder of the
 539        data in the string.  Note that this is also a logical place to call
 540        inflateEnd, but the old behaviour of only calling it on flush() is
 541        preserved.
 542     */
 543     if (err == Z_STREAM_END) {
 544         Py_XDECREF(self->unused_data);  /* Free original empty string */
 545         self->unused_data = PyString_FromStringAndSize(
 546             (char *)self->zst.next_in, self->zst.avail_in);
 547         if (self->unused_data == NULL) {
 548             Py_DECREF(RetVal);
 549             goto error;
 550         }
 551         /* We will only get Z_BUF_ERROR if the output buffer was full
 552            but there wasn't more output when we tried again, so it is
 553            not an error condition.
 554         */
 555     } else if (err != Z_OK && err != Z_BUF_ERROR) {
 556         zlib_error(self->zst, err, "while decompressing");
 557         Py_DECREF(RetVal);
 558         RetVal = NULL;
 559         goto error;
 560     }
 561
 562     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 563
 564  error:
 565     LEAVE_ZLIB
 566
 567     return RetVal;
 568 }
 569
 570 PyDoc_STRVAR(comp_flush__doc__,
 571 "flush( [mode] ) -- Return a string containing any remaining compressed data.\n"
 572 "\n"
 573 "mode can be one of the constants Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH; the\n"
 574 "default value used when mode is not specified is Z_FINISH.\n"
 575 "If mode == Z_FINISH, the compressor object can no longer be used after\n"
 576 "calling the flush() method.  Otherwise, more data can still be compressed.");
 577
 578 static PyObject *
 579 PyZlib_flush(compobject *self, PyObject *args)
 580 {
 581     int err, length = DEFAULTALLOC;
 582     PyObject *RetVal;
 583     int flushmode = Z_FINISH;
 584     unsigned long start_total_out;
 585
 586     if (!PyArg_ParseTuple(args, "|i:flush", &flushmode))
 587         return NULL;
 588
 589     /* Flushing with Z_NO_FLUSH is a no-op, so there's no point in
 590        doing any work at all; just return an empty string. */
 591     if (flushmode == Z_NO_FLUSH) {
 592         return PyString_FromStringAndSize(NULL, 0);
 593     }
 594
 595     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 596         return NULL;
 597
 598     ENTER_ZLIB
 599
 600     start_total_out = self->zst.total_out;
 601     self->zst.avail_in = 0;
 602     self->zst.avail_out = length;
 603     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 604
 605     Py_BEGIN_ALLOW_THREADS
 606     err = deflate(&(self->zst), flushmode);
 607     Py_END_ALLOW_THREADS
 608
 609     /* while Z_OK and the output buffer is full, there might be more output,
 610        so extend the output buffer and try again */
 611     while (err == Z_OK && self->zst.avail_out == 0) {
 612         if (_PyString_Resize(&RetVal, length << 1) < 0)
 613             goto error;
 614         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 615             + length;
 616         self->zst.avail_out = length;
 617         length = length << 1;
 618
 619         Py_BEGIN_ALLOW_THREADS
 620         err = deflate(&(self->zst), flushmode);
 621         Py_END_ALLOW_THREADS
 622     }
 623
 624     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 625        various data structures. Note we should only get Z_STREAM_END when
 626        flushmode is Z_FINISH, but checking both for safety*/
 627     if (err == Z_STREAM_END && flushmode == Z_FINISH) {
 628         err = deflateEnd(&(self->zst));
 629         if (err != Z_OK) {
 630             zlib_error(self->zst, err, "from deflateEnd()");
 631             Py_DECREF(RetVal);
 632             RetVal = NULL;
 633             goto error;
 634         }
 635         else
 636             self->is_initialised = 0;
 637
 638         /* We will only get Z_BUF_ERROR if the output buffer was full
 639            but there wasn't more output when we tried again, so it is
 640            not an error condition.
 641         */
 642     } else if (err!=Z_OK && err!=Z_BUF_ERROR) {
 643         zlib_error(self->zst, err, "while flushing");
 644         Py_DECREF(RetVal);
 645         RetVal = NULL;
 646         goto error;
 647     }
 648
 649     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 650
 651  error:
 652     LEAVE_ZLIB
 653
 654     return RetVal;
 655 }
 656
 657 #ifdef HAVE_ZLIB_COPY
 658 PyDoc_STRVAR(comp_copy__doc__,
 659 "copy() -- Return a copy of the compression object.");
 660
 661 static PyObject *
 662 PyZlib_copy(compobject *self)
 663 {
 664     compobject *retval = NULL;
 665     int err;
 666
 667     retval = newcompobject(&Comptype);
 668     if (!retval) return NULL;
 669
 670     /* Copy the zstream state
 671      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 672      */
 673     ENTER_ZLIB
 674     err = deflateCopy(&retval->zst, &self->zst);
 675     switch(err) {
 676     case(Z_OK):
 677         break;
 678     case(Z_STREAM_ERROR):
 679         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 680         goto error;
 681     case(Z_MEM_ERROR):
 682         PyErr_SetString(PyExc_MemoryError,
 683                         "Can't allocate memory for compression object");
 684         goto error;
 685     default:
 686         zlib_error(self->zst, err, "while copying compression object");
 687         goto error;
 688     }
 689
 690     Py_INCREF(self->unused_data);
 691     Py_INCREF(self->unconsumed_tail);
 692     Py_XDECREF(retval->unused_data);
 693     Py_XDECREF(retval->unconsumed_tail);
 694     retval->unused_data = self->unused_data;
 695     retval->unconsumed_tail = self->unconsumed_tail;
 696
 697     /* Mark it as being initialized */
 698     retval->is_initialised = 1;
 699
 700     LEAVE_ZLIB
 701     return (PyObject *)retval;
 702
 703 error:
 704     LEAVE_ZLIB
 705     Py_XDECREF(retval);
 706     return NULL;
 707 }
 708
 709 PyDoc_STRVAR(decomp_copy__doc__,
 710 "copy() -- Return a copy of the decompression object.");
 711
 712 static PyObject *
 713 PyZlib_uncopy(compobject *self)
 714 {
 715     compobject *retval = NULL;
 716     int err;
 717
 718     retval = newcompobject(&Decomptype);
 719     if (!retval) return NULL;
 720
 721     /* Copy the zstream state
 722      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 723      */
 724     ENTER_ZLIB
 725     err = inflateCopy(&retval->zst, &self->zst);
 726     switch(err) {
 727     case(Z_OK):
 728         break;
 729     case(Z_STREAM_ERROR):
 730         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 731         goto error;
 732     case(Z_MEM_ERROR):
 733         PyErr_SetString(PyExc_MemoryError,
 734                         "Can't allocate memory for decompression object");
 735         goto error;
 736     default:
 737         zlib_error(self->zst, err, "while copying decompression object");
 738         goto error;
 739     }
 740
 741     Py_INCREF(self->unused_data);
 742     Py_INCREF(self->unconsumed_tail);
 743     Py_XDECREF(retval->unused_data);
 744     Py_XDECREF(retval->unconsumed_tail);
 745     retval->unused_data = self->unused_data;
 746     retval->unconsumed_tail = self->unconsumed_tail;
 747
 748     /* Mark it as being initialized */
 749     retval->is_initialised = 1;
 750
 751     LEAVE_ZLIB
 752     return (PyObject *)retval;
 753
 754 error:
 755     LEAVE_ZLIB
 756     Py_XDECREF(retval);
 757     return NULL;
 758 }
 759 #endif
 760
 761 PyDoc_STRVAR(decomp_flush__doc__,
 762 "flush( [length] ) -- Return a string containing any remaining\n"
 763 "decompressed data. length, if given, is the initial size of the\n"
 764 "output buffer.\n"
 765 "\n"
 766 "The decompressor object can no longer be used after this call.");
 767
 768 static PyObject *
 769 PyZlib_unflush(compobject *self, PyObject *args)
 770 {
 771     int err, length = DEFAULTALLOC;
 772     PyObject * retval = NULL;
 773     unsigned long start_total_out;
 774
 775     if (!PyArg_ParseTuple(args, "|i:flush", &length))
 776         return NULL;
 777     if (!(retval = PyString_FromStringAndSize(NULL, length)))
 778         return NULL;
 779
 780
 781     ENTER_ZLIB
 782
 783     start_total_out = self->zst.total_out;
 784     self->zst.avail_out = length;
 785     self->zst.next_out = (Byte *)PyString_AS_STRING(retval);
 786
 787     Py_BEGIN_ALLOW_THREADS
 788     err = inflate(&(self->zst), Z_FINISH);
 789     Py_END_ALLOW_THREADS
 790
 791     /* while Z_OK and the output buffer is full, there might be more output,
 792        so extend the output buffer and try again */
 793     while ((err == Z_OK || err == Z_BUF_ERROR) && self->zst.avail_out == 0) {
 794         if (_PyString_Resize(&retval, length << 1) < 0)
 795             goto error;
 796         self->zst.next_out = (Byte *)PyString_AS_STRING(retval) + length;
 797         self->zst.avail_out = length;
 798         length = length << 1;
 799
 800         Py_BEGIN_ALLOW_THREADS
 801         err = inflate(&(self->zst), Z_FINISH);
 802         Py_END_ALLOW_THREADS
 803     }
 804
 805     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 806        various data structures. Note we should only get Z_STREAM_END when
 807        flushmode is Z_FINISH */
 808     if (err == Z_STREAM_END) {
 809         err = inflateEnd(&(self->zst));
 810         self->is_initialised = 0;
 811         if (err != Z_OK) {
 812             zlib_error(self->zst, err, "from inflateEnd()");
 813             Py_DECREF(retval);
 814             retval = NULL;
 815             goto error;
 816         }
 817     }
 818     _PyString_Resize(&retval, self->zst.total_out - start_total_out);
 819
 820 error:
 821
 822     LEAVE_ZLIB
 823
 824     return retval;
 825 }
 826
 827 static PyMethodDef comp_methods[] =
 828 {
 829     {"compress", (binaryfunc)PyZlib_objcompress, METH_VARARGS,
 830                  comp_compress__doc__},
 831     {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
 832               comp_flush__doc__},
 833 #ifdef HAVE_ZLIB_COPY
 834     {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
 835               comp_copy__doc__},
 836 #endif
 837     {NULL, NULL}
 838 };
 839
 840 static PyMethodDef Decomp_methods[] =
 841 {
 842     {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
 843                    decomp_decompress__doc__},
 844     {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
 845               decomp_flush__doc__},
 846 #ifdef HAVE_ZLIB_COPY
 847     {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
 848               decomp_copy__doc__},
 849 #endif
 850     {NULL, NULL}
 851 };
 852
 853 static PyObject *
 854 Comp_getattr(compobject *self, char *name)
 855 {
 856   /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
 857      internal data. */
 858
 859   return Py_FindMethod(comp_methods, (PyObject *)self, name);
 860 }
 861
 862 static PyObject *
 863 Decomp_getattr(compobject *self, char *name)
 864 {
 865     PyObject * retval;
 866
 867     ENTER_ZLIB
 868
 869     if (strcmp(name, "unused_data") == 0) {
 870         Py_INCREF(self->unused_data);
 871         retval = self->unused_data;
 872     } else if (strcmp(name, "unconsumed_tail") == 0) {
 873         Py_INCREF(self->unconsumed_tail);
 874         retval = self->unconsumed_tail;
 875     } else
 876         retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);
 877
 878     LEAVE_ZLIB
 879
 880     return retval;
 881 }
 882
 883 PyDoc_STRVAR(adler32__doc__,
 884 "adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
 885 "\n"
 886 "An optional starting value can be specified.  The returned checksum is\n"
 887 "an integer.");
 888
 889 static PyObject *
 890 PyZlib_adler32(PyObject *self, PyObject *args)
 891 {
 892     uLong adler32val = adler32(0L, Z_NULL, 0);
 893     Byte *buf;
 894     int len;
 895
 896     if (!PyArg_ParseTuple(args, "s#|k:adler32", &buf, &len, &adler32val))
 897         return NULL;
 898     adler32val = adler32(adler32val, buf, len);
 899     return PyInt_FromLong(adler32val);
 900 }
 901
 902 PyDoc_STRVAR(crc32__doc__,
 903 "crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
 904 "\n"
 905 "An optional starting value can be specified.  The returned checksum is\n"
 906 "an integer.");
 907
 908 static PyObject *
 909 PyZlib_crc32(PyObject *self, PyObject *args)
 910 {
 911     uLong crc32val = crc32(0L, Z_NULL, 0);
 912     Byte *buf;
 913     int len;
 914     if (!PyArg_ParseTuple(args, "s#|k:crc32", &buf, &len, &crc32val))
 915         return NULL;
 916     crc32val = crc32(crc32val, buf, len);
 917     return PyInt_FromLong(crc32val);
 918 }
 919
 920
 921 static PyMethodDef zlib_methods[] =
 922 {
 923     {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
 924                 adler32__doc__},
 925     {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
 926                  compress__doc__},
 927     {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
 928                     compressobj__doc__},
 929     {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
 930               crc32__doc__},
 931     {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
 932                    decompress__doc__},
 933     {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
 934                    decompressobj__doc__},
 935     {NULL, NULL}
 936 };
 937
 938 static PyTypeObject Comptype = {
 939     PyVarObject_HEAD_INIT(0, 0)
 940     "zlib.Compress",
 941     sizeof(compobject),
 942     0,
 943     (destructor)Comp_dealloc,       /*tp_dealloc*/
 944     0,                              /*tp_print*/
 945     (getattrfunc)Comp_getattr,      /*tp_getattr*/
 946     0,                              /*tp_setattr*/
 947     0,                              /*tp_compare*/
 948     0,                              /*tp_repr*/
 949     0,                              /*tp_as_number*/
 950     0,                              /*tp_as_sequence*/
 951     0,                              /*tp_as_mapping*/
 952 };
 953
 954 static PyTypeObject Decomptype = {
 955     PyVarObject_HEAD_INIT(0, 0)
 956     "zlib.Decompress",
 957     sizeof(compobject),
 958     0,
 959     (destructor)Decomp_dealloc,     /*tp_dealloc*/
 960     0,                              /*tp_print*/
 961     (getattrfunc)Decomp_getattr,    /*tp_getattr*/
 962     0,                              /*tp_setattr*/
 963     0,                              /*tp_compare*/
 964     0,                              /*tp_repr*/
 965     0,                              /*tp_as_number*/
 966     0,                              /*tp_as_sequence*/
 967     0,                              /*tp_as_mapping*/
 968 };
 969
 970 PyDoc_STRVAR(zlib_module_documentation,
 971 "The functions in this module allow compression and decompression using the\n"
 972 "zlib library, which is based on GNU zip.\n"
 973 "\n"
 974 "adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
 975 "compress(string[, level]) -- Compress string, with compression level in 1-9.\n"
 976 "compressobj([level]) -- Return a compressor object.\n"
 977 "crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
 978 "decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
 979 "decompressobj([wbits]) -- Return a decompressor object.\n"
 980 "\n"
 981 "'wbits' is window buffer size.\n"
 982 "Compressor objects support compress() and flush() methods; decompressor\n"
 983 "objects support decompress() and flush().");
 984
 985 PyMODINIT_FUNC
 986 PyInit_zlib(void)
 987 {
 988     PyObject *m, *ver;
 989     Py_TYPE(&Comptype) = &PyType_Type;
 990     Py_TYPE(&Decomptype) = &PyType_Type;
 991     m = Py_InitModule4("zlib", zlib_methods,
 992                        zlib_module_documentation,
 993                        (PyObject*)NULL,PYTHON_API_VERSION);
 994     if (m == NULL)
 995         return;
 996
 997     ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
 998     if (ZlibError != NULL) {
 999         Py_INCREF(ZlibError);
1000         PyModule_AddObject(m, "error", ZlibError);
1001     }
1002     PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
1003     PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
1004     PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
1005     PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
1006     PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
1007     PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
1008     PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
1009     PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
1010     PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
1011
1012     PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
1013     PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
1014     PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
1015     PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
1016
1017     ver = PyString_FromString(ZLIB_VERSION);
1018     if (ver != NULL)
1019         PyModule_AddObject(m, "ZLIB_VERSION", ver);
1020
1021     PyModule_AddStringConstant(m, "__version__", "1.0");
1022
1023 #ifdef WITH_THREAD
1024     zlib_lock = PyThread_allocate_lock();
1025 #endif /* WITH_THREAD */
1026 }