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