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, r_strlen=DEFAULTALLOC;
 201     z_stream zst;
 202
 203     if (!PyArg_ParseTuple(args, "s#|ii:decompress",
 204                           &input, &length, &wsize, &r_strlen))
 205         return NULL;
 206
 207     if (r_strlen <= 0)
 208         r_strlen = 1;
 209
 210     zst.avail_in = length;
 211     zst.avail_out = r_strlen;
 212
 213     if (!(result_str = PyString_FromStringAndSize(NULL, r_strlen)))
 214         return NULL;
 215
 216     zst.zalloc = (alloc_func)NULL;
 217     zst.zfree = (free_func)Z_NULL;
 218     zst.next_out = (Byte *)PyString_AS_STRING(result_str);
 219     zst.next_in = (Byte *)input;
 220     err = inflateInit2(&zst, wsize);
 221
 222     switch(err) {
 223     case(Z_OK):
 224         break;
 225     case(Z_MEM_ERROR):
 226         PyErr_SetString(PyExc_MemoryError,
 227                         "Out of memory while decompressing data");
 228         goto error;
 229     default:
 230         inflateEnd(&zst);
 231         zlib_error(zst, err, "while preparing to decompress data");
 232         goto error;
 233     }
 234
 235     do {
 236         Py_BEGIN_ALLOW_THREADS
 237         err=inflate(&zst, Z_FINISH);
 238         Py_END_ALLOW_THREADS
 239
 240         switch(err) {
 241         case(Z_STREAM_END):
 242             break;
 243         case(Z_BUF_ERROR):
 244             /*
 245              * If there is at least 1 byte of room according to zst.avail_out
 246              * and we get this error, assume that it means zlib cannot
 247              * process the inflate call() due to an error in the data.
 248              */
 249             if (zst.avail_out > 0) {
 250                 PyErr_Format(ZlibError, "Error %i while decompressing data",
 251                              err);
 252                 inflateEnd(&zst);
 253                 goto error;
 254             }
 255             /* fall through */
 256         case(Z_OK):
 257             /* need more memory */
 258             if (_PyString_Resize(&result_str, r_strlen << 1) < 0) {
 259                 inflateEnd(&zst);
 260                 goto error;
 261             }
 262             zst.next_out = (unsigned char *)PyString_AS_STRING(result_str) \
 263                 + r_strlen;
 264             zst.avail_out = r_strlen;
 265             r_strlen = r_strlen << 1;
 266             break;
 267         default:
 268             inflateEnd(&zst);
 269             zlib_error(zst, err, "while decompressing data");
 270             goto error;
 271         }
 272     } while (err != Z_STREAM_END);
 273
 274     err = inflateEnd(&zst);
 275     if (err != Z_OK) {
 276         zlib_error(zst, err, "while finishing data decompression");
 277         goto error;
 278     }
 279
 280     _PyString_Resize(&result_str, zst.total_out);
 281     return result_str;
 282
 283  error:
 284     Py_XDECREF(result_str);
 285     return NULL;
 286 }
 287
 288 static PyObject *
 289 PyZlib_compressobj(PyObject *selfptr, PyObject *args)
 290 {
 291     compobject *self;
 292     int level=Z_DEFAULT_COMPRESSION, method=DEFLATED;
 293     int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err;
 294
 295     if (!PyArg_ParseTuple(args, "|iiiii:compressobj", &level, &method, &wbits,
 296                           &memLevel, &strategy))
 297         return NULL;
 298
 299     self = newcompobject(&Comptype);
 300     if (self==NULL)
 301         return(NULL);
 302     self->zst.zalloc = (alloc_func)NULL;
 303     self->zst.zfree = (free_func)Z_NULL;
 304     self->zst.next_in = NULL;
 305     self->zst.avail_in = 0;
 306     err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
 307     switch(err) {
 308     case (Z_OK):
 309         self->is_initialised = 1;
 310         return (PyObject*)self;
 311     case (Z_MEM_ERROR):
 312         Py_DECREF(self);
 313         PyErr_SetString(PyExc_MemoryError,
 314                         "Can't allocate memory for compression object");
 315         return NULL;
 316     case(Z_STREAM_ERROR):
 317         Py_DECREF(self);
 318         PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
 319         return NULL;
 320     default:
 321         zlib_error(self->zst, err, "while creating compression object");
 322         Py_DECREF(self);
 323         return NULL;
 324     }
 325 }
 326
 327 static PyObject *
 328 PyZlib_decompressobj(PyObject *selfptr, PyObject *args)
 329 {
 330     int wbits=DEF_WBITS, err;
 331     compobject *self;
 332     if (!PyArg_ParseTuple(args, "|i:decompressobj", &wbits))
 333         return NULL;
 334
 335     self = newcompobject(&Decomptype);
 336     if (self == NULL)
 337         return(NULL);
 338     self->zst.zalloc = (alloc_func)NULL;
 339     self->zst.zfree = (free_func)Z_NULL;
 340     self->zst.next_in = NULL;
 341     self->zst.avail_in = 0;
 342     err = inflateInit2(&self->zst, wbits);
 343     switch(err) {
 344     case (Z_OK):
 345         self->is_initialised = 1;
 346         return (PyObject*)self;
 347     case(Z_STREAM_ERROR):
 348         Py_DECREF(self);
 349         PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
 350         return NULL;
 351     case (Z_MEM_ERROR):
 352         Py_DECREF(self);
 353         PyErr_SetString(PyExc_MemoryError,
 354                         "Can't allocate memory for decompression object");
 355         return NULL;
 356     default:
 357         zlib_error(self->zst, err, "while creating decompression object");
 358         Py_DECREF(self);
 359         return NULL;
 360     }
 361 }
 362
 363 static void
 364 Comp_dealloc(compobject *self)
 365 {
 366     if (self->is_initialised)
 367         deflateEnd(&self->zst);
 368     Py_XDECREF(self->unused_data);
 369     Py_XDECREF(self->unconsumed_tail);
 370     PyObject_Del(self);
 371 }
 372
 373 static void
 374 Decomp_dealloc(compobject *self)
 375 {
 376     if (self->is_initialised)
 377         inflateEnd(&self->zst);
 378     Py_XDECREF(self->unused_data);
 379     Py_XDECREF(self->unconsumed_tail);
 380     PyObject_Del(self);
 381 }
 382
 383 PyDoc_STRVAR(comp_compress__doc__,
 384 "compress(data) -- Return a string containing data compressed.\n"
 385 "\n"
 386 "After calling this function, some of the input data may still\n"
 387 "be stored in internal buffers for later processing.\n"
 388 "Call the flush() method to clear these buffers.");
 389
 390
 391 static PyObject *
 392 PyZlib_objcompress(compobject *self, PyObject *args)
 393 {
 394     int err, inplen, length = DEFAULTALLOC;
 395     PyObject *RetVal;
 396     Byte *input;
 397     unsigned long start_total_out;
 398
 399     if (!PyArg_ParseTuple(args, "s#:compress", &input, &inplen))
 400         return NULL;
 401
 402     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 403         return NULL;
 404
 405     ENTER_ZLIB
 406
 407     start_total_out = self->zst.total_out;
 408     self->zst.avail_in = inplen;
 409     self->zst.next_in = input;
 410     self->zst.avail_out = length;
 411     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 412
 413     Py_BEGIN_ALLOW_THREADS
 414     err = deflate(&(self->zst), Z_NO_FLUSH);
 415     Py_END_ALLOW_THREADS
 416
 417     /* while Z_OK and the output buffer is full, there might be more output,
 418        so extend the output buffer and try again */
 419     while (err == Z_OK && self->zst.avail_out == 0) {
 420         if (_PyString_Resize(&RetVal, length << 1) < 0)
 421             goto error;
 422         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 423             + length;
 424         self->zst.avail_out = length;
 425         length = length << 1;
 426
 427         Py_BEGIN_ALLOW_THREADS
 428         err = deflate(&(self->zst), Z_NO_FLUSH);
 429         Py_END_ALLOW_THREADS
 430     }
 431     /* We will only get Z_BUF_ERROR if the output buffer was full but
 432        there wasn't more output when we tried again, so it is not an error
 433        condition.
 434     */
 435
 436     if (err != Z_OK && err != Z_BUF_ERROR) {
 437         zlib_error(self->zst, err, "while compressing");
 438         Py_DECREF(RetVal);
 439         RetVal = NULL;
 440         goto error;
 441     }
 442     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 443
 444  error:
 445     LEAVE_ZLIB
 446     return RetVal;
 447 }
 448
 449 PyDoc_STRVAR(decomp_decompress__doc__,
 450 "decompress(data, max_length) -- Return a string containing the decompressed\n"
 451 "version of the data.\n"
 452 "\n"
 453 "After calling this function, some of the input data may still be stored in\n"
 454 "internal buffers for later processing.\n"
 455 "Call the flush() method to clear these buffers.\n"
 456 "If the max_length parameter is specified then the return value will be\n"
 457 "no longer than max_length.  Unconsumed input data will be stored in\n"
 458 "the unconsumed_tail attribute.");
 459
 460 static PyObject *
 461 PyZlib_objdecompress(compobject *self, PyObject *args)
 462 {
 463     int err, inplen, old_length, length = DEFAULTALLOC;
 464     int max_length = 0;
 465     PyObject *RetVal;
 466     Byte *input;
 467     unsigned long start_total_out;
 468
 469     if (!PyArg_ParseTuple(args, "s#|i:decompress", &input,
 470                           &inplen, &max_length))
 471         return NULL;
 472     if (max_length < 0) {
 473         PyErr_SetString(PyExc_ValueError,
 474                         "max_length must be greater than zero");
 475         return NULL;
 476     }
 477
 478     /* limit amount of data allocated to max_length */
 479     if (max_length && length > max_length)
 480         length = max_length;
 481     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 482         return NULL;
 483
 484     ENTER_ZLIB
 485
 486     start_total_out = self->zst.total_out;
 487     self->zst.avail_in = inplen;
 488     self->zst.next_in = input;
 489     self->zst.avail_out = length;
 490     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 491
 492     Py_BEGIN_ALLOW_THREADS
 493     err = inflate(&(self->zst), Z_SYNC_FLUSH);
 494     Py_END_ALLOW_THREADS
 495
 496     /* While Z_OK and the output buffer is full, there might be more output.
 497        So extend the output buffer and try again.
 498     */
 499     while (err == Z_OK && self->zst.avail_out == 0) {
 500         /* If max_length set, don't continue decompressing if we've already
 501            reached the limit.
 502         */
 503         if (max_length && length >= max_length)
 504             break;
 505
 506         /* otherwise, ... */
 507         old_length = length;
 508         length = length << 1;
 509         if (max_length && length > max_length)
 510             length = max_length;
 511
 512         if (_PyString_Resize(&RetVal, length) < 0)
 513             goto error;
 514         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 515             + old_length;
 516         self->zst.avail_out = length - old_length;
 517
 518         Py_BEGIN_ALLOW_THREADS
 519         err = inflate(&(self->zst), Z_SYNC_FLUSH);
 520         Py_END_ALLOW_THREADS
 521     }
 522
 523     /* Not all of the compressed data could be accommodated in the output buffer
 524        of specified size. Return the unconsumed tail in an attribute.*/
 525     if(max_length) {
 526         Py_DECREF(self->unconsumed_tail);
 527         self->unconsumed_tail = PyString_FromStringAndSize((char *)self->zst.next_in,
 528                                                            self->zst.avail_in);
 529         if(!self->unconsumed_tail) {
 530             Py_DECREF(RetVal);
 531             RetVal = NULL;
 532             goto error;
 533         }
 534     }
 535
 536     /* The end of the compressed data has been reached, so set the
 537        unused_data attribute to a string containing the remainder of the
 538        data in the string.  Note that this is also a logical place to call
 539        inflateEnd, but the old behaviour of only calling it on flush() is
 540        preserved.
 541     */
 542     if (err == Z_STREAM_END) {
 543         Py_XDECREF(self->unused_data);  /* Free original empty string */
 544         self->unused_data = PyString_FromStringAndSize(
 545             (char *)self->zst.next_in, self->zst.avail_in);
 546         if (self->unused_data == NULL) {
 547             Py_DECREF(RetVal);
 548             goto error;
 549         }
 550         /* We will only get Z_BUF_ERROR if the output buffer was full
 551            but there wasn't more output when we tried again, so it is
 552            not an error condition.
 553         */
 554     } else if (err != Z_OK && err != Z_BUF_ERROR) {
 555         zlib_error(self->zst, err, "while decompressing");
 556         Py_DECREF(RetVal);
 557         RetVal = NULL;
 558         goto error;
 559     }
 560
 561     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 562
 563  error:
 564     LEAVE_ZLIB
 565
 566     return RetVal;
 567 }
 568
 569 PyDoc_STRVAR(comp_flush__doc__,
 570 "flush( [mode] ) -- Return a string containing any remaining compressed data.\n"
 571 "\n"
 572 "mode can be one of the constants Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH; the\n"
 573 "default value used when mode is not specified is Z_FINISH.\n"
 574 "If mode == Z_FINISH, the compressor object can no longer be used after\n"
 575 "calling the flush() method.  Otherwise, more data can still be compressed.");
 576
 577 static PyObject *
 578 PyZlib_flush(compobject *self, PyObject *args)
 579 {
 580     int err, length = DEFAULTALLOC;
 581     PyObject *RetVal;
 582     int flushmode = Z_FINISH;
 583     unsigned long start_total_out;
 584
 585     if (!PyArg_ParseTuple(args, "|i:flush", &flushmode))
 586         return NULL;
 587
 588     /* Flushing with Z_NO_FLUSH is a no-op, so there's no point in
 589        doing any work at all; just return an empty string. */
 590     if (flushmode == Z_NO_FLUSH) {
 591         return PyString_FromStringAndSize(NULL, 0);
 592     }
 593
 594     if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
 595         return NULL;
 596
 597     ENTER_ZLIB
 598
 599     start_total_out = self->zst.total_out;
 600     self->zst.avail_in = 0;
 601     self->zst.avail_out = length;
 602     self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
 603
 604     Py_BEGIN_ALLOW_THREADS
 605     err = deflate(&(self->zst), flushmode);
 606     Py_END_ALLOW_THREADS
 607
 608     /* while Z_OK and the output buffer is full, there might be more output,
 609        so extend the output buffer and try again */
 610     while (err == Z_OK && self->zst.avail_out == 0) {
 611         if (_PyString_Resize(&RetVal, length << 1) < 0)
 612             goto error;
 613         self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
 614             + length;
 615         self->zst.avail_out = length;
 616         length = length << 1;
 617
 618         Py_BEGIN_ALLOW_THREADS
 619         err = deflate(&(self->zst), flushmode);
 620         Py_END_ALLOW_THREADS
 621     }
 622
 623     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 624        various data structures. Note we should only get Z_STREAM_END when
 625        flushmode is Z_FINISH, but checking both for safety*/
 626     if (err == Z_STREAM_END && flushmode == Z_FINISH) {
 627         err = deflateEnd(&(self->zst));
 628         if (err != Z_OK) {
 629             zlib_error(self->zst, err, "from deflateEnd()");
 630             Py_DECREF(RetVal);
 631             RetVal = NULL;
 632             goto error;
 633         }
 634         else
 635             self->is_initialised = 0;
 636
 637         /* We will only get Z_BUF_ERROR if the output buffer was full
 638            but there wasn't more output when we tried again, so it is
 639            not an error condition.
 640         */
 641     } else if (err!=Z_OK && err!=Z_BUF_ERROR) {
 642         zlib_error(self->zst, err, "while flushing");
 643         Py_DECREF(RetVal);
 644         RetVal = NULL;
 645         goto error;
 646     }
 647
 648     _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
 649
 650  error:
 651     LEAVE_ZLIB
 652
 653     return RetVal;
 654 }
 655
 656 #ifdef HAVE_ZLIB_COPY
 657 PyDoc_STRVAR(comp_copy__doc__,
 658 "copy() -- Return a copy of the compression object.");
 659
 660 static PyObject *
 661 PyZlib_copy(compobject *self)
 662 {
 663     compobject *retval = NULL;
 664     int err;
 665
 666     retval = newcompobject(&Comptype);
 667     if (!retval) return NULL;
 668
 669     /* Copy the zstream state
 670      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 671      */
 672     ENTER_ZLIB
 673     err = deflateCopy(&retval->zst, &self->zst);
 674     switch(err) {
 675     case(Z_OK):
 676         break;
 677     case(Z_STREAM_ERROR):
 678         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 679         goto error;
 680     case(Z_MEM_ERROR):
 681         PyErr_SetString(PyExc_MemoryError,
 682                         "Can't allocate memory for compression object");
 683         goto error;
 684     default:
 685         zlib_error(self->zst, err, "while copying compression object");
 686         goto error;
 687     }
 688
 689     Py_INCREF(self->unused_data);
 690     Py_INCREF(self->unconsumed_tail);
 691     Py_XDECREF(retval->unused_data);
 692     Py_XDECREF(retval->unconsumed_tail);
 693     retval->unused_data = self->unused_data;
 694     retval->unconsumed_tail = self->unconsumed_tail;
 695
 696     /* Mark it as being initialized */
 697     retval->is_initialised = 1;
 698
 699     LEAVE_ZLIB
 700     return (PyObject *)retval;
 701
 702 error:
 703     LEAVE_ZLIB
 704     Py_XDECREF(retval);
 705     return NULL;
 706 }
 707
 708 PyDoc_STRVAR(decomp_copy__doc__,
 709 "copy() -- Return a copy of the decompression object.");
 710
 711 static PyObject *
 712 PyZlib_uncopy(compobject *self)
 713 {
 714     compobject *retval = NULL;
 715     int err;
 716
 717     retval = newcompobject(&Decomptype);
 718     if (!retval) return NULL;
 719
 720     /* Copy the zstream state
 721      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 722      */
 723     ENTER_ZLIB
 724     err = inflateCopy(&retval->zst, &self->zst);
 725     switch(err) {
 726     case(Z_OK):
 727         break;
 728     case(Z_STREAM_ERROR):
 729         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 730         goto error;
 731     case(Z_MEM_ERROR):
 732         PyErr_SetString(PyExc_MemoryError,
 733                         "Can't allocate memory for decompression object");
 734         goto error;
 735     default:
 736         zlib_error(self->zst, err, "while copying decompression object");
 737         goto error;
 738     }
 739
 740     Py_INCREF(self->unused_data);
 741     Py_INCREF(self->unconsumed_tail);
 742     Py_XDECREF(retval->unused_data);
 743     Py_XDECREF(retval->unconsumed_tail);
 744     retval->unused_data = self->unused_data;
 745     retval->unconsumed_tail = self->unconsumed_tail;
 746
 747     /* Mark it as being initialized */
 748     retval->is_initialised = 1;
 749
 750     LEAVE_ZLIB
 751     return (PyObject *)retval;
 752
 753 error:
 754     LEAVE_ZLIB
 755     Py_XDECREF(retval);
 756     return NULL;
 757 }
 758 #endif
 759
 760 PyDoc_STRVAR(decomp_flush__doc__,
 761 "flush( [length] ) -- Return a string containing any remaining\n"
 762 "decompressed data. length, if given, is the initial size of the\n"
 763 "output buffer.\n"
 764 "\n"
 765 "The decompressor object can no longer be used after this call.");
 766
 767 static PyObject *
 768 PyZlib_unflush(compobject *self, PyObject *args)
 769 {
 770     int err, length = DEFAULTALLOC;
 771     PyObject * retval = NULL;
 772     unsigned long start_total_out;
 773
 774     if (!PyArg_ParseTuple(args, "|i:flush", &length))
 775         return NULL;
 776     if (!(retval = PyString_FromStringAndSize(NULL, length)))
 777         return NULL;
 778
 779
 780     ENTER_ZLIB
 781
 782     start_total_out = self->zst.total_out;
 783     self->zst.avail_out = length;
 784     self->zst.next_out = (Byte *)PyString_AS_STRING(retval);
 785
 786     Py_BEGIN_ALLOW_THREADS
 787     err = inflate(&(self->zst), Z_FINISH);
 788     Py_END_ALLOW_THREADS
 789
 790     /* while Z_OK and the output buffer is full, there might be more output,
 791        so extend the output buffer and try again */
 792     while ((err == Z_OK || err == Z_BUF_ERROR) && self->zst.avail_out == 0) {
 793         if (_PyString_Resize(&retval, length << 1) < 0)
 794             goto error;
 795         self->zst.next_out = (Byte *)PyString_AS_STRING(retval) + length;
 796         self->zst.avail_out = length;
 797         length = length << 1;
 798
 799         Py_BEGIN_ALLOW_THREADS
 800         err = inflate(&(self->zst), Z_FINISH);
 801         Py_END_ALLOW_THREADS
 802     }
 803
 804     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 805        various data structures. Note we should only get Z_STREAM_END when
 806        flushmode is Z_FINISH */
 807     if (err == Z_STREAM_END) {
 808         err = inflateEnd(&(self->zst));
 809         self->is_initialised = 0;
 810         if (err != Z_OK) {
 811             zlib_error(self->zst, err, "from inflateEnd()");
 812             Py_DECREF(retval);
 813             retval = NULL;
 814             goto error;
 815         }
 816     }
 817     _PyString_Resize(&retval, self->zst.total_out - start_total_out);
 818
 819 error:
 820
 821     LEAVE_ZLIB
 822
 823     return retval;
 824 }
 825
 826 static PyMethodDef comp_methods[] =
 827 {
 828     {"compress", (binaryfunc)PyZlib_objcompress, METH_VARARGS,
 829                  comp_compress__doc__},
 830     {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
 831               comp_flush__doc__},
 832 #ifdef HAVE_ZLIB_COPY
 833     {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
 834               comp_copy__doc__},
 835 #endif
 836     {NULL, NULL}
 837 };
 838
 839 static PyMethodDef Decomp_methods[] =
 840 {
 841     {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
 842                    decomp_decompress__doc__},
 843     {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
 844               decomp_flush__doc__},
 845 #ifdef HAVE_ZLIB_COPY
 846     {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
 847               decomp_copy__doc__},
 848 #endif
 849     {NULL, NULL}
 850 };
 851
 852 static PyObject *
 853 Comp_getattr(compobject *self, char *name)
 854 {
 855   /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
 856      internal data. */
 857
 858   return Py_FindMethod(comp_methods, (PyObject *)self, name);
 859 }
 860
 861 static PyObject *
 862 Decomp_getattr(compobject *self, char *name)
 863 {
 864     PyObject * retval;
 865
 866     ENTER_ZLIB
 867
 868     if (strcmp(name, "unused_data") == 0) {
 869         Py_INCREF(self->unused_data);
 870         retval = self->unused_data;
 871     } else if (strcmp(name, "unconsumed_tail") == 0) {
 872         Py_INCREF(self->unconsumed_tail);
 873         retval = self->unconsumed_tail;
 874     } else
 875         retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);
 876
 877     LEAVE_ZLIB
 878
 879     return retval;
 880 }
 881
 882 PyDoc_STRVAR(adler32__doc__,
 883 "adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
 884 "\n"
 885 "An optional starting value can be specified.  The returned checksum is\n"
 886 "an integer.");
 887
 888 static PyObject *
 889 PyZlib_adler32(PyObject *self, PyObject *args)
 890 {
 891     uLong adler32val = adler32(0L, Z_NULL, 0);
 892     Byte *buf;
 893     int len;
 894
 895     if (!PyArg_ParseTuple(args, "s#|k:adler32", &buf, &len, &adler32val))
 896         return NULL;
 897     adler32val = adler32(adler32val, buf, len);
 898     return PyInt_FromLong(adler32val);
 899 }
 900
 901 PyDoc_STRVAR(crc32__doc__,
 902 "crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
 903 "\n"
 904 "An optional starting value can be specified.  The returned checksum is\n"
 905 "an integer.");
 906
 907 static PyObject *
 908 PyZlib_crc32(PyObject *self, PyObject *args)
 909 {
 910     uLong crc32val = crc32(0L, Z_NULL, 0);
 911     Byte *buf;
 912     int len;
 913     if (!PyArg_ParseTuple(args, "s#|k:crc32", &buf, &len, &crc32val))
 914         return NULL;
 915     crc32val = crc32(crc32val, buf, len);
 916     return PyInt_FromLong(crc32val);
 917 }
 918
 919
 920 static PyMethodDef zlib_methods[] =
 921 {
 922     {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
 923                 adler32__doc__},
 924     {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
 925                  compress__doc__},
 926     {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
 927                     compressobj__doc__},
 928     {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
 929               crc32__doc__},
 930     {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
 931                    decompress__doc__},
 932     {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
 933                    decompressobj__doc__},
 934     {NULL, NULL}
 935 };
 936
 937 static PyTypeObject Comptype = {
 938     PyObject_HEAD_INIT(0)
 939     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     PyObject_HEAD_INIT(0)
 956     0,
 957     "zlib.Decompress",
 958     sizeof(compobject),
 959     0,
 960     (destructor)Decomp_dealloc,     /*tp_dealloc*/
 961     0,                              /*tp_print*/
 962     (getattrfunc)Decomp_getattr,    /*tp_getattr*/
 963     0,                              /*tp_setattr*/
 964     0,                              /*tp_compare*/
 965     0,                              /*tp_repr*/
 966     0,                              /*tp_as_number*/
 967     0,                              /*tp_as_sequence*/
 968     0,                              /*tp_as_mapping*/
 969 };
 970
 971 PyDoc_STRVAR(zlib_module_documentation,
 972 "The functions in this module allow compression and decompression using the\n"
 973 "zlib library, which is based on GNU zip.\n"
 974 "\n"
 975 "adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
 976 "compress(string[, level]) -- Compress string, with compression level in 1-9.\n"
 977 "compressobj([level]) -- Return a compressor object.\n"
 978 "crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
 979 "decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
 980 "decompressobj([wbits]) -- Return a decompressor object.\n"
 981 "\n"
 982 "'wbits' is window buffer size.\n"
 983 "Compressor objects support compress() and flush() methods; decompressor\n"
 984 "objects support decompress() and flush().");
 985
 986 PyMODINIT_FUNC
 987 PyInit_zlib(void)
 988 {
 989     PyObject *m, *ver;
 990     Comptype.ob_type = &PyType_Type;
 991     Decomptype.ob_type = &PyType_Type;
 992     m = Py_InitModule4("zlib", zlib_methods,
 993                        zlib_module_documentation,
 994                        (PyObject*)NULL,PYTHON_API_VERSION);
 995     if (m == NULL)
 996         return;
 997
 998     ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
 999     if (ZlibError != NULL) {
1000         Py_INCREF(ZlibError);
1001         PyModule_AddObject(m, "error", ZlibError);
1002     }
1003     PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
1004     PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
1005     PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
1006     PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
1007     PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
1008     PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
1009     PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
1010     PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
1011     PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
1012
1013     PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
1014     PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
1015     PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
1016     PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
1017
1018     ver = PyString_FromString(ZLIB_VERSION);
1019     if (ver != NULL)
1020         PyModule_AddObject(m, "ZLIB_VERSION", ver);
1021
1022     PyModule_AddStringConstant(m, "__version__", "1.0");
1023
1024 #ifdef WITH_THREAD
1025     zlib_lock = PyThread_allocate_lock();
1026 #endif /* WITH_THREAD */
1027 }