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