zlibmodule.c 25.2 KB
Newer Older
1
/* zlibmodule.c -- gzip-compatible data compression */
2 3
/* See http://www.cdrom.com/pub/infozip/zlib/ */
/* See http://www.winimage.com/zLibDll for Windows */
4

5
#include "Python.h"
Guido van Rossum's avatar
Guido van Rossum committed
6 7 8
#ifdef MS_WIN32
#define ZLIB_DLL
#endif
9
#include "zlib.h"
10 11 12 13 14 15 16 17 18 19 20

/* The following parameters are copied from zutil.h, version 0.95 */
#define DEFLATED   8
#if MAX_MEM_LEVEL >= 8
#  define DEF_MEM_LEVEL 8
#else
#  define DEF_MEM_LEVEL  MAX_MEM_LEVEL
#endif
#define DEF_WBITS MAX_WBITS

/* The output buffer will be increased in chunks of ADDCHUNK bytes. */
Jeremy Hylton's avatar
Jeremy Hylton committed
21
#define DEFAULTALLOC 16*1024
22 23 24 25 26 27 28 29 30 31 32
#define PyInit_zlib initzlib

staticforward PyTypeObject Comptype;
staticforward PyTypeObject Decomptype;

static PyObject *ZlibError;

typedef struct 
{
  PyObject_HEAD
  z_stream zst;
33
  PyObject *unused_data;
34
  int is_initialised;
35 36
} compobject;

Guido van Rossum's avatar
Guido van Rossum committed
37 38 39 40 41 42 43 44 45 46
static char compressobj__doc__[] = 
"compressobj() -- Return a compressor object.\n"
"compressobj(level) -- Return a compressor object, using the given compression level.\n"
;

static char decompressobj__doc__[] = 
"decompressobj() -- Return a decompressor object.\n"
"decompressobj(wbits) -- Return a decompressor object, setting the window buffer size to wbits.\n"
;

47 48 49 50 51 52 53 54
static compobject *
newcompobject(type)
     PyTypeObject *type;
{
        compobject *self;
        self = PyObject_NEW(compobject, type);
        if (self == NULL)
                return NULL;
55
	self->is_initialised = 0;
56
	self->unused_data = PyString_FromString("");
57 58 59
        return self;
}

Guido van Rossum's avatar
Guido van Rossum committed
60 61 62 63 64 65 66
static char compress__doc__[] = 
"compress(string) -- Compress string using the default compression level, "
"returning a string containing compressed data.\n"
"compress(string, level) -- Compress string, using the chosen compression "
"level (from 1 to 9).  Return a string containing the compressed data.\n"
;

67 68 69 70 71 72 73 74 75 76 77 78
static PyObject *
PyZlib_compress(self, args)
        PyObject *self;
        PyObject *args;
{
  PyObject *ReturnVal;
  Byte *input, *output;
  int length, level=Z_DEFAULT_COMPRESSION, err;
  z_stream zst;
  
  if (!PyArg_ParseTuple(args, "s#|i", &input, &length, &level))
    return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
79
  zst.avail_out = length + length/1000 + 12 + 1;
80 81 82 83 84 85 86
  output=(Byte*)malloc(zst.avail_out);
  if (output==NULL) 
    {
      PyErr_SetString(PyExc_MemoryError,
                      "Can't allocate memory to compress data");
      return NULL;
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
87

88
  zst.zalloc=(alloc_func)NULL;
89
  zst.zfree=(free_func)Z_NULL;
90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
  zst.next_out=(Byte *)output;
  zst.next_in =(Byte *)input;
  zst.avail_in=length;
  err=deflateInit(&zst, level);
  switch(err) 
    {
    case(Z_OK):
      break;
    case(Z_MEM_ERROR):
      PyErr_SetString(PyExc_MemoryError,
                      "Out of memory while compressing data");
      free(output);
      return NULL;
    case(Z_STREAM_ERROR):
      PyErr_SetString(ZlibError,
                      "Bad compression level");
      free(output);
      return NULL;
    default:
      {
Jeremy Hylton's avatar
Jeremy Hylton committed
110 111 112 113 114 115
	if (zst.msg == Z_NULL)
	    PyErr_Format(ZlibError, "Error %i while compressing data",
			 err); 
	else
	    PyErr_Format(ZlibError, "Error %i while compressing data: %.200s",
			 err, zst.msg);  
116 117 118 119 120 121 122 123
        deflateEnd(&zst);
        free(output);
        return NULL;
      }
    }
 
  err=deflate(&zst, Z_FINISH);
  switch(err)
Jeremy Hylton's avatar
Jeremy Hylton committed
124
  {
125 126 127 128
    case(Z_STREAM_END):
      break;
      /* Are there other errors to be trapped here? */
    default: 
Jeremy Hylton's avatar
Jeremy Hylton committed
129 130 131 132 133 134 135 136 137 138 139 140
    {
	if (zst.msg == Z_NULL)
	    PyErr_Format(ZlibError, "Error %i while compressing data",
			 err); 
	else
	    PyErr_Format(ZlibError, "Error %i while compressing data: %.200s",
			 err, zst.msg);  
	deflateEnd(&zst);
	free(output);
	return NULL;
    }
  }
141 142
  err=deflateEnd(&zst);
  if (err!=Z_OK) 
Jeremy Hylton's avatar
Jeremy Hylton committed
143 144 145 146 147 148 149 150
  {
      if (zst.msg == Z_NULL)
	  PyErr_Format(ZlibError, "Error %i while finishing compression",
		       err); 
      else
	  PyErr_Format(ZlibError,
		       "Error %i while finishing compression: %.200s",
		       err, zst.msg);  
151 152 153
      free(output);
      return NULL;
    }
154
  ReturnVal=PyString_FromStringAndSize((char *)output, zst.total_out);
155 156 157 158
  free(output);
  return ReturnVal;
}

Guido van Rossum's avatar
Guido van Rossum committed
159
static char decompress__doc__[] = 
Jeremy Hylton's avatar
Jeremy Hylton committed
160 161 162
"decompress(string) -- Decompress the data in string, returning a string containing the decompressed data.\n"
"decompress(string, wbits) -- Decompress the data in string with a window buffer size of wbits.\n"
"decompress(string, wbits, bufsize) -- Decompress the data in string with a window buffer size of wbits and an initial output buffer size of bufsize.\n"
Guido van Rossum's avatar
Guido van Rossum committed
163 164
;

165 166 167 168 169
static PyObject *
PyZlib_decompress(self, args)
        PyObject *self;
        PyObject *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
170 171
  PyObject *result_str;
  Byte *input;
172
  int length, err;
Jeremy Hylton's avatar
Jeremy Hylton committed
173
  int wsize=DEF_WBITS, r_strlen=DEFAULTALLOC;
174
  z_stream zst;
Jeremy Hylton's avatar
Jeremy Hylton committed
175
  if (!PyArg_ParseTuple(args, "s#|ii", &input, &length, &wsize, &r_strlen))
176
    return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
177 178 179

  if (r_strlen <= 0)
      r_strlen = 1;
Jeremy Hylton's avatar
Jeremy Hylton committed
180

181
  zst.avail_in=length;
Jeremy Hylton's avatar
Jeremy Hylton committed
182 183 184
  zst.avail_out=r_strlen;
  if (!(result_str = PyString_FromStringAndSize(NULL, r_strlen)))
  {
185 186 187
      PyErr_SetString(PyExc_MemoryError,
                      "Can't allocate memory to decompress data");
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
188
  }
189 190
  zst.zalloc=(alloc_func)NULL;
  zst.zfree=(free_func)Z_NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
191
  zst.next_out=(Byte *)PyString_AsString(result_str);
192
  zst.next_in =(Byte *)input;
Jeremy Hylton's avatar
Jeremy Hylton committed
193
  err=inflateInit2(&zst, wsize);
194 195 196 197 198 199 200
  switch(err)
    {
    case(Z_OK):
      break;
    case(Z_MEM_ERROR):      
      PyErr_SetString(PyExc_MemoryError,
                      "Out of memory while decompressing data");
Jeremy Hylton's avatar
Jeremy Hylton committed
201
      Py_DECREF(result_str);
202 203 204
      return NULL;
    default:
      {
Jeremy Hylton's avatar
Jeremy Hylton committed
205 206 207 208 209 210 211
	if (zst.msg == Z_NULL)
	    PyErr_Format(ZlibError, "Error %i preparing to decompress data",
			 err); 
	else
	    PyErr_Format(ZlibError,
			 "Error %i while preparing to decompress data: %.200s",
			 err, zst.msg);  
212
        inflateEnd(&zst);
Jeremy Hylton's avatar
Jeremy Hylton committed
213
	Py_DECREF(result_str);
214 215 216 217 218 219 220 221 222
        return NULL;
      }
    }
  do 
    {
      err=inflate(&zst, Z_FINISH);
      switch(err) 
        {
        case(Z_STREAM_END):
Jeremy Hylton's avatar
Jeremy Hylton committed
223
	    break;
224
	case(Z_BUF_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
225 226 227
        case(Z_OK):
	    /* need more memory */
	    if (_PyString_Resize(&result_str, r_strlen << 1) == -1)
228 229 230 231 232 233
            {
              PyErr_SetString(PyExc_MemoryError,
                              "Out of memory while decompressing data");
              inflateEnd(&zst);
              return NULL;
            }
234
	    zst.next_out = (unsigned char *)PyString_AsString(result_str) + r_strlen;
Jeremy Hylton's avatar
Jeremy Hylton committed
235 236 237
	    zst.avail_out=r_strlen;
	    r_strlen = r_strlen << 1;
	    break;
238 239
        default:
          {
Jeremy Hylton's avatar
Jeremy Hylton committed
240 241 242 243 244 245 246
	      if (zst.msg == Z_NULL)
		  PyErr_Format(ZlibError, "Error %i while decompressing data",
			       err); 
	      else
		  PyErr_Format(ZlibError,
			       "Error %i while decompressing data: %.200s",
			       err, zst.msg);  
247
            inflateEnd(&zst);
Jeremy Hylton's avatar
Jeremy Hylton committed
248
	    Py_DECREF(result_str);
249 250 251 252 253 254 255
            return NULL;
          }
        }
    } while(err!=Z_STREAM_END);
  
  err=inflateEnd(&zst);
  if (err!=Z_OK) 
Jeremy Hylton's avatar
Jeremy Hylton committed
256 257 258 259 260 261 262 263 264 265
  {
      if (zst.msg == Z_NULL)
	  PyErr_Format(ZlibError,
		       "Error %i while finishing data decompression",
		       err); 
      else
	  PyErr_Format(ZlibError,
		       "Error %i while finishing data decompression: %.200s",
		       err, zst.msg);  
      Py_DECREF(result_str);
266 267
      return NULL;
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
268 269
  _PyString_Resize(&result_str, zst.total_out);
  return result_str;
270 271 272 273 274 275 276 277 278 279
}

static PyObject *
PyZlib_compressobj(selfptr, args)
        PyObject *selfptr;
        PyObject *args;
{
  compobject *self;
  int level=Z_DEFAULT_COMPRESSION, method=DEFLATED;
  int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err;
Jeremy Hylton's avatar
Jeremy Hylton committed
280 281 282 283 284

  if (!PyArg_ParseTuple(args, "|iiiii", &level, &method, &wbits,
			&memLevel, &strategy))
      return NULL;

Jeremy Hylton's avatar
Jeremy Hylton committed
285
  self = newcompobject(&Comptype);
286
  if (self==NULL) return(NULL);
Jeremy Hylton's avatar
Jeremy Hylton committed
287 288 289
  self->zst.zalloc = (alloc_func)NULL;
  self->zst.zfree = (free_func)Z_NULL;
  err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
290 291 292
  switch(err)
    {
    case (Z_OK):
293
      self->is_initialised = 1;
294 295
      return (PyObject*)self;
    case (Z_MEM_ERROR):
296
      Py_DECREF(self);
297 298 299 300
      PyErr_SetString(PyExc_MemoryError,
                      "Can't allocate memory for compression object");
      return NULL;
    case(Z_STREAM_ERROR):
301
      Py_DECREF(self);
302
      PyErr_SetString(PyExc_ValueError,
Jeremy Hylton's avatar
Jeremy Hylton committed
303
                      "Invalid initialization option");
304 305 306
      return NULL;
    default:
      {
Jeremy Hylton's avatar
Jeremy Hylton committed
307 308 309 310 311 312 313 314
	if (self->zst.msg == Z_NULL)
	    PyErr_Format(ZlibError,
			 "Error %i while creating compression object",
			 err); 
	else
	    PyErr_Format(ZlibError,
			 "Error %i while creating compression object: %.200s",
			 err, self->zst.msg);  
315
        Py_DECREF(self);
316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333
	return NULL;
      }
    }
}

static PyObject *
PyZlib_decompressobj(selfptr, args)
        PyObject *selfptr;
        PyObject *args;
{
  int wbits=DEF_WBITS, err;
  compobject *self;
  if (!PyArg_ParseTuple(args, "|i", &wbits))
    {
     return NULL;
    }  
  self=newcompobject(&Decomptype);
  if (self==NULL) return(NULL);
334 335
  self->zst.zalloc=(alloc_func)NULL;
  self->zst.zfree=(free_func)Z_NULL;
336 337
  err=inflateInit2(&self->zst, wbits);
  switch(err)
Jeremy Hylton's avatar
Jeremy Hylton committed
338
  {
339
    case (Z_OK):
340
      self->is_initialised = 1;
341
      return (PyObject*)self;
Jeremy Hylton's avatar
Jeremy Hylton committed
342
    case(Z_STREAM_ERROR):
343
      Py_DECREF(self);
Jeremy Hylton's avatar
Jeremy Hylton committed
344 345 346
      PyErr_SetString(PyExc_ValueError,
                      "Invalid initialization option");
      return NULL;
347
    case (Z_MEM_ERROR):
348
      Py_DECREF(self);
349 350 351 352
      PyErr_SetString(PyExc_MemoryError,
                      "Can't allocate memory for decompression object");
      return NULL;
    default:
Jeremy Hylton's avatar
Jeremy Hylton committed
353 354 355 356 357 358 359 360 361
    {
	if (self->zst.msg == Z_NULL)
	    PyErr_Format(ZlibError,
			 "Error %i while creating decompression object",
			 err); 
	else
	    PyErr_Format(ZlibError,
			 "Error %i while creating decompression object: %.200s",
			 err, self->zst.msg);  
362
        Py_DECREF(self);
363 364
	return NULL;
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
365
  }
366 367 368 369 370 371
}

static void
Comp_dealloc(self)
        compobject *self;
{
372 373
    if (self->is_initialised)
      deflateEnd(&self->zst);
374
    Py_XDECREF(self->unused_data);
Jeremy Hylton's avatar
Jeremy Hylton committed
375
    PyMem_DEL(self);
376 377 378 379 380 381
}

static void
Decomp_dealloc(self)
        compobject *self;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
382
    inflateEnd(&self->zst);
383
    Py_XDECREF(self->unused_data);
Jeremy Hylton's avatar
Jeremy Hylton committed
384
    PyMem_DEL(self);
385 386
}

Guido van Rossum's avatar
Guido van Rossum committed
387 388 389 390 391 392 393 394
static char comp_compress__doc__[] =
"compress(data) -- Return a string containing a compressed version of the data.\n\n"
"After calling this function, some of the input data may still\n"
"be stored in internal buffers for later processing.\n"
"Call the flush() method to clear these buffers."
;


395 396 397 398 399
static PyObject *
PyZlib_objcompress(self, args)
        compobject *self;
        PyObject *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
400 401
  int err = Z_OK, inplen;
  int length = DEFAULTALLOC;
402 403
  PyObject *RetVal;
  Byte *input;
Jeremy Hylton's avatar
Jeremy Hylton committed
404
  unsigned long start_total_out;
405 406
  
  if (!PyArg_ParseTuple(args, "s#", &input, &inplen))
Jeremy Hylton's avatar
Jeremy Hylton committed
407
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
408 409
  self->zst.avail_in = inplen;
  self->zst.next_in = input;
Jeremy Hylton's avatar
Jeremy Hylton committed
410 411 412 413 414
  if (!(RetVal = PyString_FromStringAndSize(NULL, length))) {
      PyErr_SetString(PyExc_MemoryError,
		      "Can't allocate memory to compress data");
      return NULL;
  }
Jeremy Hylton's avatar
Jeremy Hylton committed
415
  start_total_out = self->zst.total_out;
416
  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal);
Jeremy Hylton's avatar
Jeremy Hylton committed
417 418 419 420 421 422 423 424 425 426
  self->zst.avail_out = length;
  while (self->zst.avail_in != 0 && err == Z_OK)
  {
      err = deflate(&(self->zst), Z_NO_FLUSH);
      if (self->zst.avail_out <= 0) {
	  if (_PyString_Resize(&RetVal, length << 1) == -1)  {
	      PyErr_SetString(PyExc_MemoryError,
			      "Can't allocate memory to compress data");
	      return NULL;
	  }
427
	  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal) + length;
Jeremy Hylton's avatar
Jeremy Hylton committed
428 429 430 431 432
	  self->zst.avail_out = length;
	  length = length << 1;
      }
  }
  if (err != Z_OK) 
Jeremy Hylton's avatar
Jeremy Hylton committed
433 434 435 436 437 438 439 440
  {
      if (self->zst.msg == Z_NULL)
	  PyErr_Format(ZlibError, "Error %i while compressing",
		       err); 
      else
	  PyErr_Format(ZlibError, "Error %i while compressing: %.200s",
		       err, self->zst.msg);  
      Py_DECREF(RetVal);
441 442
      return NULL;
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
443
  _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
444 445 446
  return RetVal;
}

Guido van Rossum's avatar
Guido van Rossum committed
447 448 449 450 451 452 453
static char decomp_decompress__doc__[] =
"decompress(data) -- Return a string containing the decompressed version of the data.\n\n"
"After calling this function, some of the input data may still\n"
"be stored in internal buffers for later processing.\n"
"Call the flush() method to clear these buffers."
;

454 455 456 457 458
static PyObject *
PyZlib_objdecompress(self, args)
        compobject *self;
        PyObject *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
459
  int length, err, inplen;
460 461
  PyObject *RetVal;
  Byte *input;
Jeremy Hylton's avatar
Jeremy Hylton committed
462 463
  unsigned long start_total_out;

464 465
  if (!PyArg_ParseTuple(args, "s#", &input, &inplen))
    return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
466
  start_total_out = self->zst.total_out;
Jeremy Hylton's avatar
Jeremy Hylton committed
467
  RetVal = PyString_FromStringAndSize(NULL, DEFAULTALLOC);
Jeremy Hylton's avatar
Jeremy Hylton committed
468 469
  self->zst.avail_in = inplen;
  self->zst.next_in = input;
Jeremy Hylton's avatar
Jeremy Hylton committed
470
  self->zst.avail_out = length = DEFAULTALLOC;
471
  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal);
Jeremy Hylton's avatar
Jeremy Hylton committed
472 473 474 475 476 477 478 479 480 481 482 483 484
  err = Z_OK;

  while (self->zst.avail_in != 0 && err == Z_OK)
  {
      err = inflate(&(self->zst), Z_NO_FLUSH);
      if (err == Z_OK && self->zst.avail_out <= 0) 
      {
	  if (_PyString_Resize(&RetVal, length << 1) == -1)
	  {
	      PyErr_SetString(PyExc_MemoryError,
			      "Can't allocate memory to compress data");
	      return NULL;
	  }
485
	  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal) + length;
Jeremy Hylton's avatar
Jeremy Hylton committed
486 487 488 489
	  self->zst.avail_out = length;
	  length = length << 1;
      }
  }
490

Jeremy Hylton's avatar
Jeremy Hylton committed
491 492 493 494 495 496 497 498 499
  if (err != Z_OK && err != Z_STREAM_END) 
  {
      if (self->zst.msg == Z_NULL)
	  PyErr_Format(ZlibError, "Error %i while decompressing",
		       err); 
      else
	  PyErr_Format(ZlibError, "Error %i while decompressing: %.200s",
		       err, self->zst.msg);  
      Py_DECREF(RetVal);
500 501
      return NULL;
    }
502 503 504 505 506 507 508 509 510 511 512 513 514

  if (err == Z_STREAM_END)
  {
      /* The end of the compressed data has been reached, so set 
         the unused_data attribute to a string containing the 
         remainder of the data in the string. */
      int pos = self->zst.next_in - input;  /* Position in the string */
      Py_XDECREF(self->unused_data);  /* Free the original, empty string */

      self->unused_data = PyString_FromStringAndSize(input+pos, inplen-pos);
      if (self->unused_data == NULL) return NULL;
  }

Jeremy Hylton's avatar
Jeremy Hylton committed
515
  _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
516 517 518
  return RetVal;
}

Guido van Rossum's avatar
Guido van Rossum committed
519
static char comp_flush__doc__[] =
Jeremy Hylton's avatar
Jeremy Hylton committed
520 521 522 523 524
"flush( [mode] ) -- Return a string containing any remaining compressed data.\n"
"mode can be one of the constants Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH; the \n"
"default value used when mode is not specified is Z_FINISH.\n"
"If mode == Z_FINISH, the compressor object can no longer be used after\n"
"calling the flush() method.  Otherwise, more data can still be compressed.\n"
Guido van Rossum's avatar
Guido van Rossum committed
525 526
;

527 528 529 530 531
static PyObject *
PyZlib_flush(self, args)
        compobject *self;
        PyObject *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
532
  int length=DEFAULTALLOC, err = Z_OK;
533
  PyObject *RetVal;
Jeremy Hylton's avatar
Jeremy Hylton committed
534 535 536 537
  int flushmode = Z_FINISH;
  unsigned long start_total_out;

  if (!PyArg_ParseTuple(args, "|i", &flushmode))
Jeremy Hylton's avatar
Jeremy Hylton committed
538
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
539 540 541 542 543 544 545 546

  /* Flushing with Z_NO_FLUSH is a no-op, so there's no point in
     doing any work at all; just return an empty string. */
  if (flushmode == Z_NO_FLUSH)
    {
      return PyString_FromStringAndSize(NULL, 0);
    }
  
Jeremy Hylton's avatar
Jeremy Hylton committed
547 548 549 550 551 552 553
  self->zst.avail_in = 0;
  self->zst.next_in = Z_NULL;
  if (!(RetVal = PyString_FromStringAndSize(NULL, length))) {
      PyErr_SetString(PyExc_MemoryError,
		      "Can't allocate memory to compress data");
      return NULL;
  }
Jeremy Hylton's avatar
Jeremy Hylton committed
554
  start_total_out = self->zst.total_out;
555
  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal);
Jeremy Hylton's avatar
Jeremy Hylton committed
556
  self->zst.avail_out = length;
Jeremy Hylton's avatar
Jeremy Hylton committed
557 558 559 560

  /* When flushing the zstream, there's no input data.  
     If zst.avail_out == 0, that means that more output space is
     needed to complete the flush operation. */ 
561 562 563 564 565 566 567 568 569 570 571 572 573 574 575 576 577
  while (1) {
      err = deflate(&(self->zst), flushmode);

      /* If the output is Z_OK, and there's still room in the output
	 buffer, then the flush is complete. */
      if ( (err == Z_OK) && self->zst.avail_out > 0) break;

      /* A nonzero return indicates some sort of error (but see 
	 the comment for the error handler below) */
      if ( err != Z_OK ) break;

      /* There's no space left for output, so increase the buffer and loop 
	 again */
      if (_PyString_Resize(&RetVal, length << 1) == -1)  {
	PyErr_SetString(PyExc_MemoryError,
			"Can't allocate memory to compress data");
	return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
578
      }
579 580 581
      self->zst.next_out = (unsigned char *)PyString_AsString(RetVal) + length;
      self->zst.avail_out = length;
      length = length << 1;
Jeremy Hylton's avatar
Jeremy Hylton committed
582
  }
Jeremy Hylton's avatar
Jeremy Hylton committed
583

584 585 586 587 588
  /* Raise an exception indicating an error.  The condition for
     detecting a error is kind of complicated; Z_OK indicates no
     error, but if the flushmode is Z_FINISH, then Z_STREAM_END is
     also not an error. */
  if (err!=Z_OK && !(flushmode == Z_FINISH && err == Z_STREAM_END) )
Jeremy Hylton's avatar
Jeremy Hylton committed
589
  {
Jeremy Hylton's avatar
Jeremy Hylton committed
590
      if (self->zst.msg == Z_NULL)
Jeremy Hylton's avatar
Jeremy Hylton committed
591
	  PyErr_Format(ZlibError, "Error %i while flushing",
Jeremy Hylton's avatar
Jeremy Hylton committed
592 593
		       err); 
      else
Jeremy Hylton's avatar
Jeremy Hylton committed
594
	  PyErr_Format(ZlibError, "Error %i while flushing: %.200s",
Jeremy Hylton's avatar
Jeremy Hylton committed
595 596
		       err, self->zst.msg);  
      Py_DECREF(RetVal);
597
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
598
  }
599 600 601 602

  /* If flushmode is Z_FINISH, we also have to call deflateEnd() to
     free various data structures */

Jeremy Hylton's avatar
Jeremy Hylton committed
603 604 605
  if (flushmode == Z_FINISH) {
    err=deflateEnd(&(self->zst));
    if (err!=Z_OK) {
Jeremy Hylton's avatar
Jeremy Hylton committed
606
      if (self->zst.msg == Z_NULL)
Jeremy Hylton's avatar
Jeremy Hylton committed
607 608
	PyErr_Format(ZlibError, "Error %i from deflateEnd()",
		     err); 
Jeremy Hylton's avatar
Jeremy Hylton committed
609
      else
Jeremy Hylton's avatar
Jeremy Hylton committed
610 611 612
	PyErr_Format(ZlibError,
		     "Error %i from deflateEnd(): %.200s",
		     err, self->zst.msg);  
Jeremy Hylton's avatar
Jeremy Hylton committed
613
      Py_DECREF(RetVal);
614
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
615
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
616
  }
Jeremy Hylton's avatar
Jeremy Hylton committed
617
  _PyString_Resize(&RetVal, self->zst.total_out - start_total_out);
618 619 620
  return RetVal;
}

Guido van Rossum's avatar
Guido van Rossum committed
621 622 623 624 625
static char decomp_flush__doc__[] =
"flush() -- Return a string containing any remaining decompressed data.  "
"The decompressor object can no longer be used after this call."
;

626 627 628 629 630 631 632 633 634
static PyObject *
PyZlib_unflush(self, args)
        compobject *self;
        PyObject *args;
{
  int length=0, err;
  PyObject *RetVal;
  
  if (!PyArg_NoArgs(args))
Jeremy Hylton's avatar
Jeremy Hylton committed
635
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
636 637 638 639 640 641
  if (!(RetVal = PyString_FromStringAndSize(NULL, DEFAULTALLOC)))
  {
      PyErr_SetString(PyExc_MemoryError,
		      "Can't allocate memory to decompress data");
      return NULL;
  }
642
  self->zst.avail_in=0;
643
  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal);
Jeremy Hylton's avatar
Jeremy Hylton committed
644 645 646 647 648 649 650 651 652 653 654 655 656 657
  length = self->zst.avail_out = DEFAULTALLOC;

  err = Z_OK;
  while (err == Z_OK)
  {
      err = inflate(&(self->zst), Z_FINISH);
      if (err == Z_OK && self->zst.avail_out == 0) 
      {
	  if (_PyString_Resize(&RetVal, length << 1) == -1)
	  {
	      PyErr_SetString(PyExc_MemoryError,
			      "Can't allocate memory to decompress data");
	      return NULL;
	  }
658
	  self->zst.next_out = (unsigned char *)PyString_AsString(RetVal) + length;
Jeremy Hylton's avatar
Jeremy Hylton committed
659 660 661 662
	  self->zst.avail_out = length;
	  length = length << 1;
      }
  }
663
  if (err!=Z_STREAM_END) 
Jeremy Hylton's avatar
Jeremy Hylton committed
664 665 666 667 668 669 670 671
  {
      if (self->zst.msg == Z_NULL)
	  PyErr_Format(ZlibError, "Error %i while decompressing",
		       err); 
      else
	  PyErr_Format(ZlibError, "Error %i while decompressing: %.200s",
		       err, self->zst.msg);  
      Py_DECREF(RetVal);
672
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
673
  }
674 675
  err=inflateEnd(&(self->zst));
  if (err!=Z_OK) 
Jeremy Hylton's avatar
Jeremy Hylton committed
676 677 678 679 680 681 682 683 684 685
  {
      if (self->zst.msg == Z_NULL)
	  PyErr_Format(ZlibError,
		       "Error %i while flushing decompression object",
		       err); 
      else
	  PyErr_Format(ZlibError,
		       "Error %i while flushing decompression object: %.200s",
		       err, self->zst.msg);  
      Py_DECREF(RetVal);
686
      return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
687
  }
Jeremy Hylton's avatar
Jeremy Hylton committed
688 689
  _PyString_Resize(&RetVal, 
		   (char *)self->zst.next_out - PyString_AsString(RetVal));
690 691 692 693 694
  return RetVal;
}

static PyMethodDef comp_methods[] =
{
695
        {"compress", (binaryfunc)PyZlib_objcompress, 1, comp_compress__doc__},
Jeremy Hylton's avatar
Jeremy Hylton committed
696
        {"flush", (binaryfunc)PyZlib_flush, 1, comp_flush__doc__},
697 698 699 700 701
        {NULL, NULL}
};

static PyMethodDef Decomp_methods[] =
{
702 703
        {"decompress", (binaryfunc)PyZlib_objdecompress, 1, decomp_decompress__doc__},
        {"flush", (binaryfunc)PyZlib_unflush, 0, decomp_flush__doc__},
704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719
        {NULL, NULL}
};

static PyObject *
Comp_getattr(self, name)
     compobject *self;
     char *name;
{
        return Py_FindMethod(comp_methods, (PyObject *)self, name);
}

static PyObject *
Decomp_getattr(self, name)
     compobject *self;
     char *name;
{
720 721 722 723 724
        if (strcmp(name, "unused_data") == 0) 
	  {  
	    Py_INCREF(self->unused_data);
	    return self->unused_data;
	  }
725 726 727
        return Py_FindMethod(Decomp_methods, (PyObject *)self, name);
}

Guido van Rossum's avatar
Guido van Rossum committed
728 729 730 731 732 733 734
static char adler32__doc__[] = 
"adler32(string) -- Compute an Adler-32 checksum of string, using "
"a default starting value, and returning an integer value.\n"
"adler32(string, value) -- Compute an Adler-32 checksum of string, using "
"the starting value provided, and returning an integer value\n"
;

735 736 737 738
static PyObject *
PyZlib_adler32(self, args)
     PyObject *self, *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
739 740 741
    uLong adler32val=adler32(0L, Z_NULL, 0);
    Byte *buf;
    int len;
742
 
Jeremy Hylton's avatar
Jeremy Hylton committed
743 744 745 746 747 748
    if (!PyArg_ParseTuple(args, "s#|l", &buf, &len, &adler32val))
    {
	return NULL;
    }
    adler32val = adler32(adler32val, buf, len);
    return PyInt_FromLong(adler32val);
749 750
}
     
Guido van Rossum's avatar
Guido van Rossum committed
751 752 753 754 755 756
static char crc32__doc__[] = 
"crc32(string) -- Compute a CRC-32 checksum of string, using "
"a default starting value, and returning an integer value.\n"
"crc32(string, value) -- Compute a CRC-32 checksum of string, using "
"the starting value provided, and returning an integer value.\n"
;
757 758 759 760 761

static PyObject *
PyZlib_crc32(self, args)
     PyObject *self, *args;
{
Jeremy Hylton's avatar
Jeremy Hylton committed
762 763 764 765 766 767 768 769 770
    uLong crc32val=crc32(0L, Z_NULL, 0);
    Byte *buf;
    int len;
    if (!PyArg_ParseTuple(args, "s#|l", &buf, &len, &crc32val))
    {
	return NULL;
    }
    crc32val = crc32(crc32val, buf, len);
    return PyInt_FromLong(crc32val);
771 772 773 774 775
}
     

static PyMethodDef zlib_methods[] =
{
Guido van Rossum's avatar
Guido van Rossum committed
776 777 778 779 780 781
	{"adler32", (PyCFunction)PyZlib_adler32, 1, adler32__doc__},	 
        {"compress", (PyCFunction)PyZlib_compress, 1, compress__doc__},
        {"compressobj", (PyCFunction)PyZlib_compressobj, 1, compressobj__doc__},
	{"crc32", (PyCFunction)PyZlib_crc32, 1, crc32__doc__},	 
        {"decompress", (PyCFunction)PyZlib_decompress, 1, decompress__doc__},
        {"decompressobj", (PyCFunction)PyZlib_decompressobj, 1, decompressobj__doc__},
782 783 784 785
        {NULL, NULL}
};

statichere PyTypeObject Comptype = {
Guido van Rossum's avatar
Guido van Rossum committed
786
        PyObject_HEAD_INIT(0)
787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802
        0,
        "Compress",
        sizeof(compobject),
        0,
        (destructor)Comp_dealloc,       /*tp_dealloc*/
        0,                              /*tp_print*/
        (getattrfunc)Comp_getattr,      /*tp_getattr*/
        0,                              /*tp_setattr*/
        0,                              /*tp_compare*/
        0,                              /*tp_repr*/
        0,                              /*tp_as_number*/
        0,                              /*tp_as_sequence*/
        0,                              /*tp_as_mapping*/
};

statichere PyTypeObject Decomptype = {
Guido van Rossum's avatar
Guido van Rossum committed
803
        PyObject_HEAD_INIT(0)
804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837 838 839 840
        0,
        "Decompress",
        sizeof(compobject),
        0,
        (destructor)Decomp_dealloc,     /*tp_dealloc*/
        0,                              /*tp_print*/
        (getattrfunc)Decomp_getattr,    /*tp_getattr*/
        0,                              /*tp_setattr*/
        0,                              /*tp_compare*/
        0,                              /*tp_repr*/
        0,                              /*tp_as_number*/
        0,                              /*tp_as_sequence*/
        0,                              /*tp_as_mapping*/
};

/* The following insint() routine was blatantly ripped off from 
   socketmodule.c */ 

/* Convenience routine to export an integer value.
   For simplicity, errors (which are unlikely anyway) are ignored. */
static void
insint(d, name, value)
     PyObject *d;
     char *name;
     int value;
{
	PyObject *v = PyInt_FromLong((long) value);
	if (v == NULL) {
		/* Don't bother reporting this error */
		PyErr_Clear();
	}
	else {
		PyDict_SetItemString(d, name, v);
		Py_DECREF(v);
	}
}

Guido van Rossum's avatar
Guido van Rossum committed
841 842 843 844 845 846 847 848 849 850
static char zlib_module_documentation[]=
"The functions in this module allow compression and decompression "
"using the zlib library, which is based on GNU zip.  \n\n"
"adler32(string) -- Compute an Adler-32 checksum.\n"
"adler32(string, start) -- Compute an Adler-32 checksum using a given starting value.\n"
"compress(string) -- Compress a string.\n"
"compress(string, level) -- Compress a string with the given level of compression (1--9).\n"
"compressobj([level]) -- Return a compressor object.\n"
"crc32(string) -- Compute a CRC-32 checksum.\n"
"crc32(string, start) -- Compute a CRC-32 checksum using a given starting value.\n"
Jeremy Hylton's avatar
Jeremy Hylton committed
851
"decompress(string,[wbites],[bufsize]) -- Decompresses a compressed string.\n"
Guido van Rossum's avatar
Guido van Rossum committed
852 853 854 855 856
"decompressobj([wbits]) -- Return a decompressor object (wbits=window buffer size).\n\n"
"Compressor objects support compress() and flush() methods; decompressor \n"
"objects support decompress() and flush()."
;

857
DL_EXPORT(void)
858 859
PyInit_zlib()
{
Jeremy Hylton's avatar
Jeremy Hylton committed
860
        PyObject *m, *d, *ver;
Guido van Rossum's avatar
Guido van Rossum committed
861 862
        Comptype.ob_type = &PyType_Type;
        Decomptype.ob_type = &PyType_Type;
Guido van Rossum's avatar
Guido van Rossum committed
863 864 865
        m = Py_InitModule4("zlib", zlib_methods,
			   zlib_module_documentation,
			   (PyObject*)NULL,PYTHON_API_VERSION);
866
        d = PyModule_GetDict(m);
867
        ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
868
        PyDict_SetItemString(d, "error", ZlibError);
Jeremy Hylton's avatar
Jeremy Hylton committed
869

870 871 872
	insint(d, "MAX_WBITS", MAX_WBITS);
	insint(d, "DEFLATED", DEFLATED);
	insint(d, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
Jeremy Hylton's avatar
Jeremy Hylton committed
873 874 875 876 877 878
	insint(d, "Z_BEST_SPEED", Z_BEST_SPEED);
	insint(d, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
	insint(d, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
	insint(d, "Z_FILTERED", Z_FILTERED);
	insint(d, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
	insint(d, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
Jeremy Hylton's avatar
Jeremy Hylton committed
879 880 881 882 883 884

	insint(d, "Z_FINISH", Z_FINISH);
	insint(d, "Z_NO_FLUSH", Z_NO_FLUSH);
	insint(d, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
	insint(d, "Z_FULL_FLUSH", Z_FULL_FLUSH);

Jeremy Hylton's avatar
Jeremy Hylton committed
885 886
	ver = PyString_FromString(ZLIB_VERSION);
	PyDict_SetItemString(d, "ZLIB_VERSION", ver);
887
	Py_DECREF(ver);
888
}