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 (length <= 0) {
 778         PyErr_SetString(PyExc_ValueError, "length must be greater than zero");
 779         return NULL;
 780     }
 781     if (!(retval = PyString_FromStringAndSize(NULL, length)))
 782         return NULL;
 783
 784
 785     ENTER_ZLIB
 786
 787     start_total_out = self->zst.total_out;
 788     self->zst.avail_out = length;
 789     self->zst.next_out = (Byte *)PyString_AS_STRING(retval);
 790
 791     Py_BEGIN_ALLOW_THREADS
 792     err = inflate(&(self->zst), Z_FINISH);
 793     Py_END_ALLOW_THREADS
 794
 795     /* while Z_OK and the output buffer is full, there might be more output,
 796        so extend the output buffer and try again */
 797     while ((err == Z_OK || err == Z_BUF_ERROR) && self->zst.avail_out == 0) {
 798         if (_PyString_Resize(&retval, length << 1) < 0)
 799             goto error;
 800         self->zst.next_out = (Byte *)PyString_AS_STRING(retval) + length;
 801         self->zst.avail_out = length;
 802         length = length << 1;
 803
 804         Py_BEGIN_ALLOW_THREADS
 805         err = inflate(&(self->zst), Z_FINISH);
 806         Py_END_ALLOW_THREADS
 807     }
 808
 809     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 810        various data structures. Note we should only get Z_STREAM_END when
 811        flushmode is Z_FINISH */
 812     if (err == Z_STREAM_END) {
 813         err = inflateEnd(&(self->zst));
 814         self->is_initialised = 0;
 815         if (err != Z_OK) {
 816             zlib_error(self->zst, err, "from inflateEnd()");
 817             Py_DECREF(retval);
 818             retval = NULL;
 819             goto error;
 820         }
 821     }
 822     _PyString_Resize(&retval, self->zst.total_out - start_total_out);
 823
 824 error:
 825
 826     LEAVE_ZLIB
 827
 828     return retval;
 829 }
 830
 831 static PyMethodDef comp_methods[] =
 832 {
 833     {"compress", (binaryfunc)PyZlib_objcompress, METH_VARARGS,
 834                  comp_compress__doc__},
 835     {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
 836               comp_flush__doc__},
 837 #ifdef HAVE_ZLIB_COPY
 838     {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
 839               comp_copy__doc__},
 840 #endif
 841     {NULL, NULL}
 842 };
 843
 844 static PyMethodDef Decomp_methods[] =
 845 {
 846     {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
 847                    decomp_decompress__doc__},
 848     {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
 849               decomp_flush__doc__},
 850 #ifdef HAVE_ZLIB_COPY
 851     {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
 852               decomp_copy__doc__},
 853 #endif
 854     {NULL, NULL}
 855 };
 856
 857 static PyObject *
 858 Comp_getattr(compobject *self, char *name)
 859 {
 860   /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
 861      internal data. */
 862
 863   return Py_FindMethod(comp_methods, (PyObject *)self, name);
 864 }
 865
 866 static PyObject *
 867 Decomp_getattr(compobject *self, char *name)
 868 {
 869     PyObject * retval;
 870
 871     ENTER_ZLIB
 872
 873     if (strcmp(name, "unused_data") == 0) {
 874         Py_INCREF(self->unused_data);
 875         retval = self->unused_data;
 876     } else if (strcmp(name, "unconsumed_tail") == 0) {
 877         Py_INCREF(self->unconsumed_tail);
 878         retval = self->unconsumed_tail;
 879     } else
 880         retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);
 881
 882     LEAVE_ZLIB
 883
 884     return retval;
 885 }
 886
 887 PyDoc_STRVAR(adler32__doc__,
 888 "adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
 889 "\n"
 890 "An optional starting value can be specified.  The returned checksum is\n"
 891 "a signed integer.");
 892
 893 static PyObject *
 894 PyZlib_adler32(PyObject *self, PyObject *args)
 895 {
 896     unsigned int adler32val = 1;  /* adler32(0L, Z_NULL, 0) */
 897     Byte *buf;
 898     int len, signed_val;
 899
 900     if (!PyArg_ParseTuple(args, "s#|I:adler32", &buf, &len, &adler32val))
 901         return NULL;
 902     /* In Python 2.x we return a signed integer regardless of native platform
 903      * long size (the 32bit unsigned long is treated as 32-bit signed and sign
 904      * extended into a 64-bit long inside the integer object).  3.0 does the
 905      * right thing and returns unsigned. http://bugs.python.org/issue1202 */
 906     signed_val = adler32(adler32val, buf, len);
 907     return PyInt_FromLong(signed_val);
 908 }
 909
 910 PyDoc_STRVAR(crc32__doc__,
 911 "crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
 912 "\n"
 913 "An optional starting value can be specified.  The returned checksum is\n"
 914 "a signed integer.");
 915
 916 static PyObject *
 917 PyZlib_crc32(PyObject *self, PyObject *args)
 918 {
 919     unsigned int crc32val = 0;  /* crc32(0L, Z_NULL, 0) */
 920     Byte *buf;
 921     int len, signed_val;
 922
 923     if (!PyArg_ParseTuple(args, "s#|I:crc32", &buf, &len, &crc32val))
 924         return NULL;
 925     /* In Python 2.x we return a signed integer regardless of native platform
 926      * long size (the 32bit unsigned long is treated as 32-bit signed and sign
 927      * extended into a 64-bit long inside the integer object).  3.0 does the
 928      * right thing and returns unsigned. http://bugs.python.org/issue1202 */
 929     signed_val = crc32(crc32val, buf, len);
 930     return PyInt_FromLong(signed_val);
 931 }
 932
 933
 934 static PyMethodDef zlib_methods[] =
 935 {
 936     {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
 937                 adler32__doc__},
 938     {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
 939                  compress__doc__},
 940     {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
 941                     compressobj__doc__},
 942     {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
 943               crc32__doc__},
 944     {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
 945                    decompress__doc__},
 946     {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
 947                    decompressobj__doc__},
 948     {NULL, NULL}
 949 };
 950
 951 static PyTypeObject Comptype = {
 952     PyVarObject_HEAD_INIT(0, 0)
 953     "zlib.Compress",
 954     sizeof(compobject),
 955     0,
 956     (destructor)Comp_dealloc,       /*tp_dealloc*/
 957     0,                              /*tp_print*/
 958     (getattrfunc)Comp_getattr,      /*tp_getattr*/
 959     0,                              /*tp_setattr*/
 960     0,                              /*tp_compare*/
 961     0,                              /*tp_repr*/
 962     0,                              /*tp_as_number*/
 963     0,                              /*tp_as_sequence*/
 964     0,                              /*tp_as_mapping*/
 965 };
 966
 967 static PyTypeObject Decomptype = {
 968     PyVarObject_HEAD_INIT(0, 0)
 969     "zlib.Decompress",
 970     sizeof(compobject),
 971     0,
 972     (destructor)Decomp_dealloc,     /*tp_dealloc*/
 973     0,                              /*tp_print*/
 974     (getattrfunc)Decomp_getattr,    /*tp_getattr*/
 975     0,                              /*tp_setattr*/
 976     0,                              /*tp_compare*/
 977     0,                              /*tp_repr*/
 978     0,                              /*tp_as_number*/
 979     0,                              /*tp_as_sequence*/
 980     0,                              /*tp_as_mapping*/
 981 };
 982
 983 PyDoc_STRVAR(zlib_module_documentation,
 984 "The functions in this module allow compression and decompression using the\n"
 985 "zlib library, which is based on GNU zip.\n"
 986 "\n"
 987 "adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
 988 "compress(string[, level]) -- Compress string, with compression level in 1-9.\n"
 989 "compressobj([level]) -- Return a compressor object.\n"
 990 "crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
 991 "decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
 992 "decompressobj([wbits]) -- Return a decompressor object.\n"
 993 "\n"
 994 "'wbits' is window buffer size.\n"
 995 "Compressor objects support compress() and flush() methods; decompressor\n"
 996 "objects support decompress() and flush().");
 997
 998 PyMODINIT_FUNC
 999 PyInit_zlib(void)
1000 {
1001     PyObject *m, *ver;
1002     Py_TYPE(&Comptype) = &PyType_Type;
1003     Py_TYPE(&Decomptype) = &PyType_Type;
1004     m = Py_InitModule4("zlib", zlib_methods,
1005                        zlib_module_documentation,
1006                        (PyObject*)NULL,PYTHON_API_VERSION);
1007     if (m == NULL)
1008         return;
1009
1010     ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
1011     if (ZlibError != NULL) {
1012         Py_INCREF(ZlibError);
1013         PyModule_AddObject(m, "error", ZlibError);
1014     }
1015     PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
1016     PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
1017     PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
1018     PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
1019     PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
1020     PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
1021     PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
1022     PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
1023     PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
1024
1025     PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
1026     PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
1027     PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
1028     PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
1029
1030     ver = PyString_FromString(ZLIB_VERSION);
1031     if (ver != NULL)
1032         PyModule_AddObject(m, "ZLIB_VERSION", ver);
1033
1034     PyModule_AddStringConstant(m, "__version__", "1.0");
1035
1036 #ifdef WITH_THREAD
1037     zlib_lock = PyThread_allocate_lock();
1038 #endif /* WITH_THREAD */
1039 }