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 PyDoc_STRVAR(comp_copy__doc__,
 657 "copy() -- Return a copy of the compression object.");
 658
 659 static PyObject *
 660 PyZlib_copy(compobject *self)
 661 {
 662     compobject *retval = NULL;
 663     int err;
 664
 665     retval = newcompobject(&Comptype);
 666     if (!retval) return NULL;
 667
 668     /* Copy the zstream state
 669      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 670      */
 671     ENTER_ZLIB
 672     err = deflateCopy(&retval->zst, &self->zst);
 673     switch(err) {
 674     case(Z_OK):
 675         break;
 676     case(Z_STREAM_ERROR):
 677         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 678         goto error;
 679     case(Z_MEM_ERROR):
 680         PyErr_SetString(PyExc_MemoryError,
 681                         "Can't allocate memory for compression object");
 682         goto error;
 683     default:
 684         zlib_error(self->zst, err, "while copying compression object");
 685         goto error;
 686     }
 687
 688     Py_INCREF(self->unused_data);
 689     Py_INCREF(self->unconsumed_tail);
 690     Py_XDECREF(retval->unused_data);
 691     Py_XDECREF(retval->unconsumed_tail);
 692     retval->unused_data = self->unused_data;
 693     retval->unconsumed_tail = self->unconsumed_tail;
 694
 695     /* Mark it as being initialized */
 696     retval->is_initialised = 1;
 697
 698     LEAVE_ZLIB
 699     return (PyObject *)retval;
 700
 701 error:
 702     LEAVE_ZLIB
 703     Py_XDECREF(retval);
 704     return NULL;
 705 }
 706
 707 PyDoc_STRVAR(decomp_copy__doc__,
 708 "copy() -- Return a copy of the decompression object.");
 709
 710 static PyObject *
 711 PyZlib_uncopy(compobject *self)
 712 {
 713     compobject *retval = NULL;
 714     int err;
 715
 716     retval = newcompobject(&Decomptype);
 717     if (!retval) return NULL;
 718
 719     /* Copy the zstream state
 720      * We use ENTER_ZLIB / LEAVE_ZLIB to make this thread-safe
 721      */
 722     ENTER_ZLIB
 723     err = inflateCopy(&retval->zst, &self->zst);
 724     switch(err) {
 725     case(Z_OK):
 726         break;
 727     case(Z_STREAM_ERROR):
 728         PyErr_SetString(PyExc_ValueError, "Inconsistent stream state");
 729         goto error;
 730     case(Z_MEM_ERROR):
 731         PyErr_SetString(PyExc_MemoryError,
 732                         "Can't allocate memory for decompression object");
 733         goto error;
 734     default:
 735         zlib_error(self->zst, err, "while copying decompression object");
 736         goto error;
 737     }
 738
 739     Py_INCREF(self->unused_data);
 740     Py_INCREF(self->unconsumed_tail);
 741     Py_XDECREF(retval->unused_data);
 742     Py_XDECREF(retval->unconsumed_tail);
 743     retval->unused_data = self->unused_data;
 744     retval->unconsumed_tail = self->unconsumed_tail;
 745
 746     /* Mark it as being initialized */
 747     retval->is_initialised = 1;
 748
 749     LEAVE_ZLIB
 750     return (PyObject *)retval;
 751
 752 error:
 753     LEAVE_ZLIB
 754     Py_XDECREF(retval);
 755     return NULL;
 756 }
 757
 758 PyDoc_STRVAR(decomp_flush__doc__,
 759 "flush( [length] ) -- Return a string containing any remaining\n"
 760 "decompressed data. length, if given, is the initial size of the\n"
 761 "output buffer.\n"
 762 "\n"
 763 "The decompressor object can no longer be used after this call.");
 764
 765 static PyObject *
 766 PyZlib_unflush(compobject *self, PyObject *args)
 767 {
 768     int err, length = DEFAULTALLOC;
 769     PyObject * retval = NULL;
 770     unsigned long start_total_out;
 771
 772     if (!PyArg_ParseTuple(args, "|i:flush", &length))
 773         return NULL;
 774     if (!(retval = PyString_FromStringAndSize(NULL, length)))
 775         return NULL;
 776
 777
 778     ENTER_ZLIB
 779
 780     start_total_out = self->zst.total_out;
 781     self->zst.avail_out = length;
 782     self->zst.next_out = (Byte *)PyString_AS_STRING(retval);
 783
 784     Py_BEGIN_ALLOW_THREADS
 785     err = inflate(&(self->zst), Z_FINISH);
 786     Py_END_ALLOW_THREADS
 787
 788     /* while Z_OK and the output buffer is full, there might be more output,
 789        so extend the output buffer and try again */
 790     while ((err == Z_OK || err == Z_BUF_ERROR) && self->zst.avail_out == 0) {
 791         if (_PyString_Resize(&retval, length << 1) < 0)
 792             goto error;
 793         self->zst.next_out = (Byte *)PyString_AS_STRING(retval) + length;
 794         self->zst.avail_out = length;
 795         length = length << 1;
 796
 797         Py_BEGIN_ALLOW_THREADS
 798         err = inflate(&(self->zst), Z_FINISH);
 799         Py_END_ALLOW_THREADS
 800     }
 801
 802     /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
 803        various data structures. Note we should only get Z_STREAM_END when
 804        flushmode is Z_FINISH */
 805     if (err == Z_STREAM_END) {
 806         err = inflateEnd(&(self->zst));
 807         self->is_initialised = 0;
 808         if (err != Z_OK) {
 809             zlib_error(self->zst, err, "from inflateEnd()");
 810             Py_DECREF(retval);
 811             retval = NULL;
 812             goto error;
 813         }
 814     }
 815     _PyString_Resize(&retval, self->zst.total_out - start_total_out);
 816
 817 error:
 818
 819     LEAVE_ZLIB
 820
 821     return retval;
 822 }
 823
 824 static PyMethodDef comp_methods[] =
 825 {
 826     {"compress", (binaryfunc)PyZlib_objcompress, METH_VARARGS,
 827                  comp_compress__doc__},
 828     {"flush", (binaryfunc)PyZlib_flush, METH_VARARGS,
 829               comp_flush__doc__},
 830     {"copy",  (PyCFunction)PyZlib_copy, METH_NOARGS,
 831               comp_copy__doc__},
 832     {NULL, NULL}
 833 };
 834
 835 static PyMethodDef Decomp_methods[] =
 836 {
 837     {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
 838                    decomp_decompress__doc__},
 839     {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
 840               decomp_flush__doc__},
 841     {"copy",  (PyCFunction)PyZlib_uncopy, METH_NOARGS,
 842               decomp_copy__doc__},
 843     {NULL, NULL}
 844 };
 845
 846 static PyObject *
 847 Comp_getattr(compobject *self, char *name)
 848 {
 849   /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
 850      internal data. */
 851
 852   return Py_FindMethod(comp_methods, (PyObject *)self, name);
 853 }
 854
 855 static PyObject *
 856 Decomp_getattr(compobject *self, char *name)
 857 {
 858     PyObject * retval;
 859
 860     ENTER_ZLIB
 861
 862     if (strcmp(name, "unused_data") == 0) {
 863         Py_INCREF(self->unused_data);
 864         retval = self->unused_data;
 865     } else if (strcmp(name, "unconsumed_tail") == 0) {
 866         Py_INCREF(self->unconsumed_tail);
 867         retval = self->unconsumed_tail;
 868     } else
 869         retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);
 870
 871     LEAVE_ZLIB
 872
 873     return retval;
 874 }
 875
 876 PyDoc_STRVAR(adler32__doc__,
 877 "adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
 878 "\n"
 879 "An optional starting value can be specified.  The returned checksum is\n"
 880 "an integer.");
 881
 882 static PyObject *
 883 PyZlib_adler32(PyObject *self, PyObject *args)
 884 {
 885     uLong adler32val = adler32(0L, Z_NULL, 0);
 886     Byte *buf;
 887     int len;
 888
 889     if (!PyArg_ParseTuple(args, "s#|k:adler32", &buf, &len, &adler32val))
 890         return NULL;
 891     adler32val = adler32(adler32val, buf, len);
 892     return PyInt_FromLong(adler32val);
 893 }
 894
 895 PyDoc_STRVAR(crc32__doc__,
 896 "crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
 897 "\n"
 898 "An optional starting value can be specified.  The returned checksum is\n"
 899 "an integer.");
 900
 901 static PyObject *
 902 PyZlib_crc32(PyObject *self, PyObject *args)
 903 {
 904     uLong crc32val = crc32(0L, Z_NULL, 0);
 905     Byte *buf;
 906     int len;
 907     if (!PyArg_ParseTuple(args, "s#|k:crc32", &buf, &len, &crc32val))
 908         return NULL;
 909     crc32val = crc32(crc32val, buf, len);
 910     return PyInt_FromLong(crc32val);
 911 }
 912
 913
 914 static PyMethodDef zlib_methods[] =
 915 {
 916     {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
 917                 adler32__doc__},
 918     {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
 919                  compress__doc__},
 920     {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
 921                     compressobj__doc__},
 922     {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
 923               crc32__doc__},
 924     {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
 925                    decompress__doc__},
 926     {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
 927                    decompressobj__doc__},
 928     {NULL, NULL}
 929 };
 930
 931 static PyTypeObject Comptype = {
 932     PyObject_HEAD_INIT(0)
 933     0,
 934     "zlib.Compress",
 935     sizeof(compobject),
 936     0,
 937     (destructor)Comp_dealloc,       /*tp_dealloc*/
 938     0,                              /*tp_print*/
 939     (getattrfunc)Comp_getattr,      /*tp_getattr*/
 940     0,                              /*tp_setattr*/
 941     0,                              /*tp_compare*/
 942     0,                              /*tp_repr*/
 943     0,                              /*tp_as_number*/
 944     0,                              /*tp_as_sequence*/
 945     0,                              /*tp_as_mapping*/
 946 };
 947
 948 static PyTypeObject Decomptype = {
 949     PyObject_HEAD_INIT(0)
 950     0,
 951     "zlib.Decompress",
 952     sizeof(compobject),
 953     0,
 954     (destructor)Decomp_dealloc,     /*tp_dealloc*/
 955     0,                              /*tp_print*/
 956     (getattrfunc)Decomp_getattr,    /*tp_getattr*/
 957     0,                              /*tp_setattr*/
 958     0,                              /*tp_compare*/
 959     0,                              /*tp_repr*/
 960     0,                              /*tp_as_number*/
 961     0,                              /*tp_as_sequence*/
 962     0,                              /*tp_as_mapping*/
 963 };
 964
 965 PyDoc_STRVAR(zlib_module_documentation,
 966 "The functions in this module allow compression and decompression using the\n"
 967 "zlib library, which is based on GNU zip.\n"
 968 "\n"
 969 "adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
 970 "compress(string[, level]) -- Compress string, with compression level in 1-9.\n"
 971 "compressobj([level]) -- Return a compressor object.\n"
 972 "crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
 973 "decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
 974 "decompressobj([wbits]) -- Return a decompressor object.\n"
 975 "\n"
 976 "'wbits' is window buffer size.\n"
 977 "Compressor objects support compress() and flush() methods; decompressor\n"
 978 "objects support decompress() and flush().");
 979
 980 PyMODINIT_FUNC
 981 PyInit_zlib(void)
 982 {
 983     PyObject *m, *ver;
 984     Comptype.ob_type = &PyType_Type;
 985     Decomptype.ob_type = &PyType_Type;
 986     m = Py_InitModule4("zlib", zlib_methods,
 987                        zlib_module_documentation,
 988                        (PyObject*)NULL,PYTHON_API_VERSION);
 989     if (m == NULL)
 990         return;
 991
 992     ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
 993     if (ZlibError != NULL) {
 994         Py_INCREF(ZlibError);
 995         PyModule_AddObject(m, "error", ZlibError);
 996     }
 997     PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
 998     PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
 999     PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
1000     PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
1001     PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
1002     PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
1003     PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
1004     PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
1005     PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
1006
1007     PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
1008     PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
1009     PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
1010     PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
1011
1012     ver = PyString_FromString(ZLIB_VERSION);
1013     if (ver != NULL)
1014         PyModule_AddObject(m, "ZLIB_VERSION", ver);
1015
1016     PyModule_AddStringConstant(m, "__version__", "1.0");
1017
1018 #ifdef WITH_THREAD
1019     zlib_lock = PyThread_allocate_lock();
1020 #endif /* WITH_THREAD */
1021 }