zlibmodule.c 24.3 KB
Newer Older
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 8
#include "Python.h"
#include "zlib.h"
9

10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
#ifdef WITH_THREAD
#include "pythread.h"

/* #defs ripped off from _tkinter.c, even though the situation here is much
   simpler, because we don't have to worry about waiting for Tcl
   events!  And, since zlib itself is threadsafe, we don't need to worry
   about re-entering zlib functions.

   N.B.

   Since ENTER_ZLIB and LEAVE_ZLIB only need to be called on functions
   that modify the components of preexisting de/compress objects, it
   could prove to be a performance gain on multiprocessor machines if
   there was an de/compress object-specific lock.  However, for the
   moment the ENTER_ZLIB and LEAVE_ZLIB calls are global for ALL
   de/compress objects.
 */

static PyThread_type_lock zlib_lock = NULL; /* initialized on module load */

#define ENTER_ZLIB \
31 32 33
	Py_BEGIN_ALLOW_THREADS \
	PyThread_acquire_lock(zlib_lock, 1); \
	Py_END_ALLOW_THREADS
34 35

#define LEAVE_ZLIB \
36
	PyThread_release_lock(zlib_lock);
37 38 39 40 41 42 43 44

#else

#define ENTER_ZLIB
#define LEAVE_ZLIB

#endif

45 46 47 48 49 50 51 52 53
/* 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

54 55
/* The output buffer will be increased in chunks of DEFAULTALLOC bytes. */
#define DEFAULTALLOC (16*1024)
56 57 58 59 60 61 62
#define PyInit_zlib initzlib

staticforward PyTypeObject Comptype;
staticforward PyTypeObject Decomptype;

static PyObject *ZlibError;

63
typedef struct
64
{
Jeremy Hylton's avatar
Jeremy Hylton committed
65 66 67 68 69
    PyObject_HEAD
    z_stream zst;
    PyObject *unused_data;
    PyObject *unconsumed_tail;
    int is_initialised;
70 71
} compobject;

Jeremy Hylton's avatar
Jeremy Hylton committed
72 73 74 75 76 77 78 79 80
static void
zlib_error(z_stream zst, int err, char *msg)
{
    if (zst.msg == Z_NULL)
	PyErr_Format(ZlibError, "Error %d %s", err, msg);
    else
	PyErr_Format(ZlibError, "Error %d %s: %.200s", err, msg, zst.msg);
}

81
static char compressobj__doc__[] =
82 83 84
"compressobj([level]) -- Return a compressor object.\n"
"\n"
"Optional arg level is the compression level, in 1-9.";
Guido van Rossum's avatar
Guido van Rossum committed
85

86
static char decompressobj__doc__[] =
87 88 89
"decompressobj([wbits]) -- Return a decompressor object.\n"
"\n"
"Optional arg wbits is the window buffer size.";
Guido van Rossum's avatar
Guido van Rossum committed
90

91
static compobject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
92
newcompobject(PyTypeObject *type)
93
{
94
    compobject *self;
Jeremy Hylton's avatar
Jeremy Hylton committed
95 96 97 98 99 100 101 102 103 104 105 106 107 108 109
    self = PyObject_New(compobject, type);
    if (self == NULL)
	return NULL;
    self->is_initialised = 0;
    self->unused_data = PyString_FromString("");
    if (self->unused_data == NULL) {
	Py_DECREF(self);
	return NULL;
    }
    self->unconsumed_tail = PyString_FromString("");
    if (self->unconsumed_tail == NULL) {
	Py_DECREF(self);
	return NULL;
    }
    return self;
110 111
}

112
static char compress__doc__[] =
113 114 115
"compress(string[, level]) -- Returned compressed string.\n"
"\n"
"Optional arg level is the compression level, in 1-9.";
Guido van Rossum's avatar
Guido van Rossum committed
116

117
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
118
PyZlib_compress(PyObject *self, PyObject *args)
119
{
Jeremy Hylton's avatar
Jeremy Hylton committed
120 121 122 123
    PyObject *ReturnVal = NULL;
    Byte *input, *output;
    int length, level=Z_DEFAULT_COMPRESSION, err;
    z_stream zst;
124

Jeremy Hylton's avatar
Jeremy Hylton committed
125
    /* require Python string object, optional 'level' arg */
126
    if (!PyArg_ParseTuple(args, "s#|i:compress", &input, &length, &level))
Jeremy Hylton's avatar
Jeremy Hylton committed
127
	return NULL;
128

Jeremy Hylton's avatar
Jeremy Hylton committed
129
    zst.avail_out = length + length/1000 + 12 + 1;
130

131 132
    output = (Byte*)malloc(zst.avail_out);
    if (output == NULL) {
Jeremy Hylton's avatar
Jeremy Hylton committed
133 134 135
	PyErr_SetString(PyExc_MemoryError,
			"Can't allocate memory to compress data");
	return NULL;
136
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
137

Jeremy Hylton's avatar
Jeremy Hylton committed
138 139
    /* Past the point of no return.  From here on out, we need to make sure
       we clean up mallocs & INCREFs. */
140

141 142 143 144 145 146
    zst.zalloc = (alloc_func)NULL;
    zst.zfree = (free_func)Z_NULL;
    zst.next_out = (Byte *)output;
    zst.next_in = (Byte *)input;
    zst.avail_in = length;
    err = deflateInit(&zst, level);
147

148
    switch(err) {
149
    case(Z_OK):
Jeremy Hylton's avatar
Jeremy Hylton committed
150
	break;
151
    case(Z_MEM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
152 153
	PyErr_SetString(PyExc_MemoryError,
			"Out of memory while compressing data");
154
	goto error;
155
    case(Z_STREAM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
156 157
	PyErr_SetString(ZlibError,
			"Bad compression level");
158
	goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
159
    default:
160
        deflateEnd(&zst);
Jeremy Hylton's avatar
Jeremy Hylton committed
161
	zlib_error(zst, err, "while compressing data");
162
	goto error;
163
    }
164

165 166 167
    Py_BEGIN_ALLOW_THREADS;
    err = deflate(&zst, Z_FINISH);
    Py_END_ALLOW_THREADS;
168

169 170 171 172
    if (err != Z_STREAM_END) {
	zlib_error(zst, err, "while compressing data");
	deflateEnd(&zst);
	goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
173
    }
174

175 176
    err=deflateEnd(&zst);
    if (err == Z_OK)
177
	ReturnVal = PyString_FromStringAndSize((char *)output,
178
					       zst.total_out);
179
    else
180 181 182
	zlib_error(zst, err, "while finishing compression");

 error:
Jeremy Hylton's avatar
Jeremy Hylton committed
183
    free(output);
184

Jeremy Hylton's avatar
Jeremy Hylton committed
185
    return ReturnVal;
186 187
}

188
static char decompress__doc__[] =
189 190 191 192
"decompress(string[, wbits[, bufsize]]) -- Return decompressed string.\n"
"\n"
"Optional arg wbits is the window buffer size.  Optional arg bufsize is\n"
"the initial output buffer size.";
Guido van Rossum's avatar
Guido van Rossum committed
193

194
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
195
PyZlib_decompress(PyObject *self, PyObject *args)
196
{
Jeremy Hylton's avatar
Jeremy Hylton committed
197 198 199 200 201 202
    PyObject *result_str;
    Byte *input;
    int length, err;
    int wsize=DEF_WBITS, r_strlen=DEFAULTALLOC;
    z_stream zst;

203
    if (!PyArg_ParseTuple(args, "s#|ii:decompress",
204
			  &input, &length, &wsize, &r_strlen))
Jeremy Hylton's avatar
Jeremy Hylton committed
205
	return NULL;
206

Jeremy Hylton's avatar
Jeremy Hylton committed
207 208
    if (r_strlen <= 0)
	r_strlen = 1;
209

Jeremy Hylton's avatar
Jeremy Hylton committed
210 211
    zst.avail_in = length;
    zst.avail_out = r_strlen;
212

213
    if (!(result_str = PyString_FromStringAndSize(NULL, r_strlen)))
Jeremy Hylton's avatar
Jeremy Hylton committed
214
	return NULL;
215

Jeremy Hylton's avatar
Jeremy Hylton committed
216 217
    zst.zalloc = (alloc_func)NULL;
    zst.zfree = (free_func)Z_NULL;
218
    zst.next_out = (Byte *)PyString_AS_STRING(result_str);
Jeremy Hylton's avatar
Jeremy Hylton committed
219 220 221 222
    zst.next_in = (Byte *)input;
    err = inflateInit2(&zst, wsize);

    switch(err) {
223
    case(Z_OK):
Jeremy Hylton's avatar
Jeremy Hylton committed
224
	break;
225
    case(Z_MEM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
226 227
	PyErr_SetString(PyExc_MemoryError,
			"Out of memory while decompressing data");
228
	goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
229
    default:
230
        inflateEnd(&zst);
Jeremy Hylton's avatar
Jeremy Hylton committed
231
	zlib_error(zst, err, "while preparing to decompress data");
232
	goto error;
233
    }
234

Jeremy Hylton's avatar
Jeremy Hylton committed
235 236 237 238
    do {
	Py_BEGIN_ALLOW_THREADS
	err=inflate(&zst, Z_FINISH);
	Py_END_ALLOW_THREADS
239

Jeremy Hylton's avatar
Jeremy Hylton committed
240 241
	switch(err) {
	case(Z_STREAM_END):
Jeremy Hylton's avatar
Jeremy Hylton committed
242
	    break;
243
	case(Z_BUF_ERROR):
244 245 246 247 248
	    /*
	     * If there is at least 1 byte of room according to zst.avail_out
	     * and we get this error, assume that it means zlib cannot
	     * process the inflate call() due to an error in the data.
	     */
Jeremy Hylton's avatar
Jeremy Hylton committed
249
	    if (zst.avail_out > 0) {
Jeremy Hylton's avatar
Jeremy Hylton committed
250 251 252
		PyErr_Format(ZlibError, "Error %i while decompressing data",
			     err);
		inflateEnd(&zst);
253
		goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
254
	    }
255 256
	    /* fall through */
	case(Z_OK):
Jeremy Hylton's avatar
Jeremy Hylton committed
257
	    /* need more memory */
Jeremy Hylton's avatar
Jeremy Hylton committed
258 259 260
	    if (_PyString_Resize(&result_str, r_strlen << 1) == -1) {
		inflateEnd(&zst);
		result_str = NULL;
261
		goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
262
	    }
263
	    zst.next_out = (unsigned char *)PyString_AS_STRING(result_str) \
Jeremy Hylton's avatar
Jeremy Hylton committed
264 265
		+ r_strlen;
	    zst.avail_out = r_strlen;
Jeremy Hylton's avatar
Jeremy Hylton committed
266 267
	    r_strlen = r_strlen << 1;
	    break;
Jeremy Hylton's avatar
Jeremy Hylton committed
268
	default:
Jeremy Hylton's avatar
Jeremy Hylton committed
269
	    inflateEnd(&zst);
Jeremy Hylton's avatar
Jeremy Hylton committed
270
	    zlib_error(zst, err, "while decompressing data");
271
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
272 273
	}
    } while (err != Z_STREAM_END);
274

275 276 277
    err = inflateEnd(&zst);
    if (err != Z_OK) {
	zlib_error(zst, err, "while finishing data decompression");
278
	goto error;
279
    }
280

281
    _PyString_Resize(&result_str, zst.total_out);
Jeremy Hylton's avatar
Jeremy Hylton committed
282
    return result_str;
283 284 285 286

 error:
    Py_XDECREF(result_str);
    return NULL;
287 288 289
}

static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
290
PyZlib_compressobj(PyObject *selfptr, PyObject *args)
291
{
Jeremy Hylton's avatar
Jeremy Hylton committed
292 293 294 295 296 297 298 299 300
    compobject *self;
    int level=Z_DEFAULT_COMPRESSION, method=DEFLATED;
    int wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=0, err;

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

    self = newcompobject(&Comptype);
301
    if (self==NULL)
Jeremy Hylton's avatar
Jeremy Hylton committed
302 303 304 305 306
	return(NULL);
    self->zst.zalloc = (alloc_func)NULL;
    self->zst.zfree = (free_func)Z_NULL;
    err = deflateInit2(&self->zst, level, method, wbits, memLevel, strategy);
    switch(err) {
307
    case (Z_OK):
Jeremy Hylton's avatar
Jeremy Hylton committed
308 309
	self->is_initialised = 1;
	return (PyObject*)self;
310
    case (Z_MEM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
311 312 313 314
	Py_DECREF(self);
	PyErr_SetString(PyExc_MemoryError,
			"Can't allocate memory for compression object");
	return NULL;
315
    case(Z_STREAM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
316
	Py_DECREF(self);
317
	PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
Jeremy Hylton's avatar
Jeremy Hylton committed
318
	return NULL;
319
    default:
Jeremy Hylton's avatar
Jeremy Hylton committed
320
	zlib_error(self->zst, err, "while creating compression object");
321
        Py_DECREF(self);
322 323 324 325 326
	return NULL;
    }
}

static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
327
PyZlib_decompressobj(PyObject *selfptr, PyObject *args)
328
{
Jeremy Hylton's avatar
Jeremy Hylton committed
329 330 331 332 333 334
    int wbits=DEF_WBITS, err;
    compobject *self;
    if (!PyArg_ParseTuple(args, "|i:decompressobj", &wbits))
	return NULL;

    self = newcompobject(&Decomptype);
335
    if (self == NULL)
Jeremy Hylton's avatar
Jeremy Hylton committed
336 337 338 339 340
	return(NULL);
    self->zst.zalloc = (alloc_func)NULL;
    self->zst.zfree = (free_func)Z_NULL;
    err = inflateInit2(&self->zst, wbits);
    switch(err) {
341
    case (Z_OK):
Jeremy Hylton's avatar
Jeremy Hylton committed
342 343
	self->is_initialised = 1;
	return (PyObject*)self;
Jeremy Hylton's avatar
Jeremy Hylton committed
344
    case(Z_STREAM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
345
	Py_DECREF(self);
346
	PyErr_SetString(PyExc_ValueError, "Invalid initialization option");
Jeremy Hylton's avatar
Jeremy Hylton committed
347
	return NULL;
348
    case (Z_MEM_ERROR):
Jeremy Hylton's avatar
Jeremy Hylton committed
349 350 351 352
	Py_DECREF(self);
	PyErr_SetString(PyExc_MemoryError,
			"Can't allocate memory for decompression object");
	return NULL;
353
    default:
Jeremy Hylton's avatar
Jeremy Hylton committed
354
	zlib_error(self->zst, err, "while creating decompression object");
355
        Py_DECREF(self);
356
	return NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
357
    }
358 359 360
}

static void
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
361
Comp_dealloc(compobject *self)
362
{
363
    if (self->is_initialised)
Jeremy Hylton's avatar
Jeremy Hylton committed
364
	deflateEnd(&self->zst);
365
    Py_XDECREF(self->unused_data);
366
    Py_XDECREF(self->unconsumed_tail);
367
    PyObject_Del(self);
368 369 370
}

static void
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
371
Decomp_dealloc(compobject *self)
372
{
373
    if (self->is_initialised)
Jeremy Hylton's avatar
Jeremy Hylton committed
374
	inflateEnd(&self->zst);
375
    Py_XDECREF(self->unused_data);
376
    Py_XDECREF(self->unconsumed_tail);
377
    PyObject_Del(self);
378 379
}

Guido van Rossum's avatar
Guido van Rossum committed
380
static char comp_compress__doc__[] =
381 382
"compress(data) -- Return a string containing data compressed.\n"
"\n"
Guido van Rossum's avatar
Guido van Rossum committed
383 384
"After calling this function, some of the input data may still\n"
"be stored in internal buffers for later processing.\n"
385
"Call the flush() method to clear these buffers.";
Guido van Rossum's avatar
Guido van Rossum committed
386 387


388
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
389
PyZlib_objcompress(compobject *self, PyObject *args)
390
{
Jeremy Hylton's avatar
Jeremy Hylton committed
391 392 393 394
    int err, inplen, length = DEFAULTALLOC;
    PyObject *RetVal;
    Byte *input;
    unsigned long start_total_out;
395

396
    if (!PyArg_ParseTuple(args, "s#:compress", &input, &inplen))
Jeremy Hylton's avatar
Jeremy Hylton committed
397 398
	return NULL;

399
    if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
Jeremy Hylton's avatar
Jeremy Hylton committed
400
	return NULL;
401

Jeremy Hylton's avatar
Jeremy Hylton committed
402
    ENTER_ZLIB
403

Jeremy Hylton's avatar
Jeremy Hylton committed
404 405 406
    start_total_out = self->zst.total_out;
    self->zst.avail_in = inplen;
    self->zst.next_in = input;
407
    self->zst.avail_out = length;
408
    self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
Jeremy Hylton's avatar
Jeremy Hylton committed
409

410
    Py_BEGIN_ALLOW_THREADS
411
    err = deflate(&(self->zst), Z_NO_FLUSH);
412
    Py_END_ALLOW_THREADS
Jeremy Hylton's avatar
Jeremy Hylton committed
413 414 415 416

    /* while Z_OK and the output buffer is full, there might be more output,
       so extend the output buffer and try again */
    while (err == Z_OK && self->zst.avail_out == 0) {
417 418 419
	if (_PyString_Resize(&RetVal, length << 1) == -1) {
	    RetVal = NULL;
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
420
	}
421
	self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
Jeremy Hylton's avatar
Jeremy Hylton committed
422 423 424
	    + length;
	self->zst.avail_out = length;
	length = length << 1;
425

Jeremy Hylton's avatar
Jeremy Hylton committed
426 427 428
	Py_BEGIN_ALLOW_THREADS
	err = deflate(&(self->zst), Z_NO_FLUSH);
	Py_END_ALLOW_THREADS
429
    }
430
    /* We will only get Z_BUF_ERROR if the output buffer was full but
Jeremy Hylton's avatar
Jeremy Hylton committed
431
       there wasn't more output when we tried again, so it is not an error
432 433
       condition.
    */
434

435 436 437 438 439
    if (err != Z_OK && err != Z_BUF_ERROR) {
	zlib_error(self->zst, err, "while compressing");
	Py_DECREF(RetVal);
	RetVal = NULL;
	goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
440
    }
441
    if (_PyString_Resize(&RetVal,
442 443
			 self->zst.total_out - start_total_out) < 0)
	RetVal = NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
444

445
 error:
Jeremy Hylton's avatar
Jeremy Hylton committed
446 447
    LEAVE_ZLIB
    return RetVal;
448 449
}

Guido van Rossum's avatar
Guido van Rossum committed
450
static char decomp_decompress__doc__[] =
451 452 453 454 455
"decompress(data, max_length) -- Return a string containing the decompressed\n"
"version of the data.\n"
"\n"
"After calling this function, some of the input data may still be stored in\n"
"internal buffers for later processing.\n"
456 457 458
"Call the flush() method to clear these buffers.\n"
"If the max_length parameter is specified then the return value will be\n"
"no longer than max_length.  Unconsumed input data will be stored in\n"
459
"the unconsumed_tail attribute.";
Guido van Rossum's avatar
Guido van Rossum committed
460

461
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
462
PyZlib_objdecompress(compobject *self, PyObject *args)
463
{
Jeremy Hylton's avatar
Jeremy Hylton committed
464 465 466 467 468 469
    int err, inplen, old_length, length = DEFAULTALLOC;
    int max_length = 0;
    PyObject *RetVal;
    Byte *input;
    unsigned long start_total_out;

470
    if (!PyArg_ParseTuple(args, "s#|i:decompress", &input,
471
			  &inplen, &max_length))
Jeremy Hylton's avatar
Jeremy Hylton committed
472 473 474 475 476 477
	return NULL;
    if (max_length < 0) {
	PyErr_SetString(PyExc_ValueError,
			"max_length must be greater than zero");
	return NULL;
    }
478

Jeremy Hylton's avatar
Jeremy Hylton committed
479
    /* limit amount of data allocated to max_length */
480
    if (max_length && length > max_length)
Jeremy Hylton's avatar
Jeremy Hylton committed
481
	length = max_length;
482
    if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
Jeremy Hylton's avatar
Jeremy Hylton committed
483
	return NULL;
484

Jeremy Hylton's avatar
Jeremy Hylton committed
485
    ENTER_ZLIB
486

Jeremy Hylton's avatar
Jeremy Hylton committed
487 488 489 490
    start_total_out = self->zst.total_out;
    self->zst.avail_in = inplen;
    self->zst.next_in = input;
    self->zst.avail_out = length;
491
    self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
492

Jeremy Hylton's avatar
Jeremy Hylton committed
493 494 495
    Py_BEGIN_ALLOW_THREADS
    err = inflate(&(self->zst), Z_SYNC_FLUSH);
    Py_END_ALLOW_THREADS
496

Jeremy Hylton's avatar
Jeremy Hylton committed
497 498
    /* While Z_OK and the output buffer is full, there might be more output.
       So extend the output buffer and try again.
499
    */
500
    while (err == Z_OK && self->zst.avail_out == 0) {
Jeremy Hylton's avatar
Jeremy Hylton committed
501 502 503 504 505
	/* If max_length set, don't continue decompressing if we've already
	   reached the limit.
	*/
	if (max_length && length >= max_length)
	    break;
506

Jeremy Hylton's avatar
Jeremy Hylton committed
507 508 509
	/* otherwise, ... */
	old_length = length;
	length = length << 1;
510
	if (max_length && length > max_length)
Jeremy Hylton's avatar
Jeremy Hylton committed
511
	    length = max_length;
512

Jeremy Hylton's avatar
Jeremy Hylton committed
513
	if (_PyString_Resize(&RetVal, length) == -1) {
514 515
	    RetVal = NULL;
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
516
	}
517
	self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
Jeremy Hylton's avatar
Jeremy Hylton committed
518 519
	    + old_length;
	self->zst.avail_out = length - old_length;
520

Jeremy Hylton's avatar
Jeremy Hylton committed
521 522 523 524
	Py_BEGIN_ALLOW_THREADS
	err = inflate(&(self->zst), Z_SYNC_FLUSH);
	Py_END_ALLOW_THREADS
    }
525

Jeremy Hylton's avatar
Jeremy Hylton committed
526 527 528 529
    /* Not all of the compressed data could be accomodated in the output buffer
       of specified size. Return the unconsumed tail in an attribute.*/
    if(max_length) {
	Py_DECREF(self->unconsumed_tail);
Jack Jansen's avatar
Jack Jansen committed
530
	self->unconsumed_tail = PyString_FromStringAndSize((char *)self->zst.next_in,
Jeremy Hylton's avatar
Jeremy Hylton committed
531
							   self->zst.avail_in);
532 533 534 535 536
	if(!self->unconsumed_tail) {
	    Py_DECREF(RetVal);
	    RetVal = NULL;
	    goto error;
	}
537
    }
Jeremy Hylton's avatar
Jeremy Hylton committed
538

539 540 541
    /* 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.  Note that this is also a logical place to call
Jeremy Hylton's avatar
Jeremy Hylton committed
542 543 544
       inflateEnd, but the old behaviour of only calling it on flush() is
       preserved.
    */
545 546 547 548 549
    if (err == Z_STREAM_END) {
	Py_XDECREF(self->unused_data);  /* Free original empty string */
	self->unused_data = PyString_FromStringAndSize(
	    (char *)self->zst.next_in, self->zst.avail_in);
	if (self->unused_data == NULL) {
Jeremy Hylton's avatar
Jeremy Hylton committed
550
	    Py_DECREF(RetVal);
551
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
552
	}
553
	/* We will only get Z_BUF_ERROR if the output buffer was full
554 555 556 557 558 559 560 561
	   but there wasn't more output when we tried again, so it is
	   not an error condition.
	*/
    } else if (err != Z_OK && err != Z_BUF_ERROR) {
	zlib_error(self->zst, err, "while decompressing");
	Py_DECREF(RetVal);
	RetVal = NULL;
	goto error;
562
    }
563

564 565
    if (_PyString_Resize(&RetVal, self->zst.total_out - start_total_out) < 0)
	RetVal = NULL;
566

567
 error:
Jeremy Hylton's avatar
Jeremy Hylton committed
568
    LEAVE_ZLIB
569

Jeremy Hylton's avatar
Jeremy Hylton committed
570
    return RetVal;
571 572
}

Guido van Rossum's avatar
Guido van Rossum committed
573
static char comp_flush__doc__[] =
Jeremy Hylton's avatar
Jeremy Hylton committed
574
"flush( [mode] ) -- Return a string containing any remaining compressed data.\n"
575 576
"\n"
"mode can be one of the constants Z_SYNC_FLUSH, Z_FULL_FLUSH, Z_FINISH; the\n"
Jeremy Hylton's avatar
Jeremy Hylton committed
577 578
"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"
579
"calling the flush() method.  Otherwise, more data can still be compressed.\n";
Guido van Rossum's avatar
Guido van Rossum committed
580

581
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
582
PyZlib_flush(compobject *self, PyObject *args)
583
{
Jeremy Hylton's avatar
Jeremy Hylton committed
584 585 586 587
    int err, length = DEFAULTALLOC;
    PyObject *RetVal;
    int flushmode = Z_FINISH;
    unsigned long start_total_out;
588

Jeremy Hylton's avatar
Jeremy Hylton committed
589 590
    if (!PyArg_ParseTuple(args, "|i:flush", &flushmode))
	return NULL;
591

Jeremy Hylton's avatar
Jeremy Hylton committed
592 593 594 595 596
    /* 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);
    }
597

598
    if (!(RetVal = PyString_FromStringAndSize(NULL, length)))
Jeremy Hylton's avatar
Jeremy Hylton committed
599 600 601
	return NULL;

    ENTER_ZLIB
602

Jeremy Hylton's avatar
Jeremy Hylton committed
603 604
    start_total_out = self->zst.total_out;
    self->zst.avail_in = 0;
605
    self->zst.avail_out = length;
606
    self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal);
607 608

    Py_BEGIN_ALLOW_THREADS
609
    err = deflate(&(self->zst), flushmode);
610 611
    Py_END_ALLOW_THREADS

Jeremy Hylton's avatar
Jeremy Hylton committed
612 613 614 615
    /* while Z_OK and the output buffer is full, there might be more output,
       so extend the output buffer and try again */
    while (err == Z_OK && self->zst.avail_out == 0) {
	if (_PyString_Resize(&RetVal, length << 1) == -1)  {
616 617
	    RetVal = NULL;
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
618
	}
619
	self->zst.next_out = (unsigned char *)PyString_AS_STRING(RetVal) \
Jeremy Hylton's avatar
Jeremy Hylton committed
620 621 622 623 624 625 626 627
	    + length;
	self->zst.avail_out = length;
	length = length << 1;

	Py_BEGIN_ALLOW_THREADS
	err = deflate(&(self->zst), flushmode);
	Py_END_ALLOW_THREADS
    }
628

Jeremy Hylton's avatar
Jeremy Hylton committed
629
    /* If flushmode is Z_FINISH, we also have to call deflateEnd() to free
630
       various data structures. Note we should only get Z_STREAM_END when
Jeremy Hylton's avatar
Jeremy Hylton committed
631
       flushmode is Z_FINISH, but checking both for safety*/
632 633 634 635
    if (err == Z_STREAM_END && flushmode == Z_FINISH) {
	err = deflateEnd(&(self->zst));
	if (err != Z_OK) {
	    zlib_error(self->zst, err, "from deflateEnd()");
Jeremy Hylton's avatar
Jeremy Hylton committed
636
	    Py_DECREF(RetVal);
637 638
	    RetVal = NULL;
	    goto error;
Jeremy Hylton's avatar
Jeremy Hylton committed
639
	}
640 641
	else
	    self->is_initialised = 0;
642 643 644

	/* We will only get Z_BUF_ERROR if the output buffer was full
	   but there wasn't more output when we tried again, so it is
645 646 647 648 649 650
	   not an error condition.
	*/
    } else if (err!=Z_OK && err!=Z_BUF_ERROR) {
	zlib_error(self->zst, err, "while flushing");
	Py_DECREF(RetVal);
	RetVal = NULL;
Jeremy Hylton's avatar
Jeremy Hylton committed
651
    }
652

653 654
    if (_PyString_Resize(&RetVal, self->zst.total_out - start_total_out) < 0)
	RetVal = NULL;
655

656
 error:
657
    LEAVE_ZLIB
658 659

    return RetVal;
660 661
}

Guido van Rossum's avatar
Guido van Rossum committed
662
static char decomp_flush__doc__[] =
663 664 665
"flush() -- Return a string containing any remaining decompressed data.\n"
"\n"
"The decompressor object can no longer be used after this call.";
Guido van Rossum's avatar
Guido van Rossum committed
666

667
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
668
PyZlib_unflush(compobject *self, PyObject *args)
669
/*decompressor flush is a no-op because all pending data would have been
670
  flushed by the decompress method. However, this routine previously called
671
  inflateEnd, causing any further decompress or flush calls to raise
672
  exceptions. This behaviour has been preserved.*/
673
{
Jeremy Hylton's avatar
Jeremy Hylton committed
674
    int err;
675
    PyObject * retval = NULL;
676

Jeremy Hylton's avatar
Jeremy Hylton committed
677 678
    if (!PyArg_ParseTuple(args, ""))
	return NULL;
679

Jeremy Hylton's avatar
Jeremy Hylton committed
680
    ENTER_ZLIB
681

Jeremy Hylton's avatar
Jeremy Hylton committed
682
    err = inflateEnd(&(self->zst));
683
    if (err != Z_OK)
Jeremy Hylton's avatar
Jeremy Hylton committed
684
	zlib_error(self->zst, err, "from inflateEnd()");
685
    else {
Jeremy Hylton's avatar
Jeremy Hylton committed
686 687 688
	self->is_initialised = 0;
	retval = PyString_FromStringAndSize(NULL, 0);
    }
689

Jeremy Hylton's avatar
Jeremy Hylton committed
690
    LEAVE_ZLIB
691

Jeremy Hylton's avatar
Jeremy Hylton committed
692
    return retval;
693 694 695 696
}

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

static PyMethodDef Decomp_methods[] =
{
706
    {"decompress", (binaryfunc)PyZlib_objdecompress, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
707
                   decomp_decompress__doc__},
708
    {"flush", (binaryfunc)PyZlib_unflush, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
709 710
              decomp_flush__doc__},
    {NULL, NULL}
711 712 713
};

static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
714
Comp_getattr(compobject *self, char *name)
715
{
716 717 718 719
  /* No ENTER/LEAVE_ZLIB is necessary because this fn doesn't touch
     internal data. */

  return Py_FindMethod(comp_methods, (PyObject *)self, name);
720 721 722
}

static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
723
Decomp_getattr(compobject *self, char *name)
724
{
Jeremy Hylton's avatar
Jeremy Hylton committed
725 726 727 728
    PyObject * retval;

    ENTER_ZLIB

729
    if (strcmp(name, "unused_data") == 0) {
Jeremy Hylton's avatar
Jeremy Hylton committed
730 731
	Py_INCREF(self->unused_data);
	retval = self->unused_data;
732
    } else if (strcmp(name, "unconsumed_tail") == 0) {
Jeremy Hylton's avatar
Jeremy Hylton committed
733 734
	Py_INCREF(self->unconsumed_tail);
	retval = self->unconsumed_tail;
735
    } else
Jeremy Hylton's avatar
Jeremy Hylton committed
736 737 738 739 740
	retval = Py_FindMethod(Decomp_methods, (PyObject *)self, name);

    LEAVE_ZLIB

    return retval;
741 742
}

743
static char adler32__doc__[] =
744 745 746 747
"adler32(string[, start]) -- Compute an Adler-32 checksum of string.\n"
"\n"
"An optional starting value can be specified.  The returned checksum is\n"
"an integer.";
Guido van Rossum's avatar
Guido van Rossum committed
748

749
static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
750
PyZlib_adler32(PyObject *self, PyObject *args)
751
{
Jeremy Hylton's avatar
Jeremy Hylton committed
752
    uLong adler32val = adler32(0L, Z_NULL, 0);
Jeremy Hylton's avatar
Jeremy Hylton committed
753 754
    Byte *buf;
    int len;
755

756
    if (!PyArg_ParseTuple(args, "s#|l:adler32", &buf, &len, &adler32val))
Jeremy Hylton's avatar
Jeremy Hylton committed
757 758 759
	return NULL;
    adler32val = adler32(adler32val, buf, len);
    return PyInt_FromLong(adler32val);
760
}
761 762

static char crc32__doc__[] =
763 764 765 766
"crc32(string[, start]) -- Compute a CRC-32 checksum of string.\n"
"\n"
"An optional starting value can be specified.  The returned checksum is\n"
"an integer.";
767 768

static PyObject *
Peter Schneider-Kamp's avatar
Peter Schneider-Kamp committed
769
PyZlib_crc32(PyObject *self, PyObject *args)
770
{
Jeremy Hylton's avatar
Jeremy Hylton committed
771
    uLong crc32val = crc32(0L, Z_NULL, 0);
Jeremy Hylton's avatar
Jeremy Hylton committed
772 773
    Byte *buf;
    int len;
774
    if (!PyArg_ParseTuple(args, "s#|l:crc32", &buf, &len, &crc32val))
Jeremy Hylton's avatar
Jeremy Hylton committed
775 776 777
	return NULL;
    crc32val = crc32(crc32val, buf, len);
    return PyInt_FromLong(crc32val);
778
}
779

780 781 782

static PyMethodDef zlib_methods[] =
{
783
    {"adler32", (PyCFunction)PyZlib_adler32, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
784
                adler32__doc__},
785
    {"compress", (PyCFunction)PyZlib_compress,  METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
786
                 compress__doc__},
787
    {"compressobj", (PyCFunction)PyZlib_compressobj, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
788
                    compressobj__doc__},
789 790 791
    {"crc32", (PyCFunction)PyZlib_crc32, METH_VARARGS,
              crc32__doc__},
    {"decompress", (PyCFunction)PyZlib_decompress, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
792
                   decompress__doc__},
793
    {"decompressobj", (PyCFunction)PyZlib_decompressobj, METH_VARARGS,
Jeremy Hylton's avatar
Jeremy Hylton committed
794 795
                   decompressobj__doc__},
    {NULL, NULL}
796 797 798
};

statichere PyTypeObject Comptype = {
Jeremy Hylton's avatar
Jeremy Hylton committed
799 800
    PyObject_HEAD_INIT(0)
    0,
801
    "zlib.Compress",
Jeremy Hylton's avatar
Jeremy Hylton committed
802 803 804 805 806 807 808 809 810 811 812
    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*/
813 814 815
};

statichere PyTypeObject Decomptype = {
Jeremy Hylton's avatar
Jeremy Hylton committed
816 817
    PyObject_HEAD_INIT(0)
    0,
818
    "zlib.Decompress",
Jeremy Hylton's avatar
Jeremy Hylton committed
819 820 821 822 823 824 825 826 827 828 829
    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*/
830 831
};

Guido van Rossum's avatar
Guido van Rossum committed
832
static char zlib_module_documentation[]=
833 834 835 836 837
"The functions in this module allow compression and decompression using the\n"
"zlib library, which is based on GNU zip.\n"
"\n"
"adler32(string[, start]) -- Compute an Adler-32 checksum.\n"
"compress(string[, level]) -- Compress string, with compression level in 1-9.\n"
Guido van Rossum's avatar
Guido van Rossum committed
838
"compressobj([level]) -- Return a compressor object.\n"
839
"crc32(string[, start]) -- Compute a CRC-32 checksum.\n"
840
"decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
841 842 843 844 845
"decompressobj([wbits]) -- Return a decompressor object.\n"
"\n"
"'wbits' is window buffer size.\n"
"Compressor objects support compress() and flush() methods; decompressor\n"
"objects support decompress() and flush().";
Guido van Rossum's avatar
Guido van Rossum committed
846

847
DL_EXPORT(void)
848
PyInit_zlib(void)
849
{
Jeremy Hylton's avatar
Jeremy Hylton committed
850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869
    PyObject *m, *d, *ver;
    Comptype.ob_type = &PyType_Type;
    Decomptype.ob_type = &PyType_Type;
    m = Py_InitModule4("zlib", zlib_methods,
		       zlib_module_documentation,
		       (PyObject*)NULL,PYTHON_API_VERSION);
    d = PyModule_GetDict(m);
    ZlibError = PyErr_NewException("zlib.error", NULL, NULL);
    if (ZlibError != NULL)
	PyDict_SetItemString(d, "error", ZlibError);

    PyModule_AddIntConstant(m, "MAX_WBITS", MAX_WBITS);
    PyModule_AddIntConstant(m, "DEFLATED", DEFLATED);
    PyModule_AddIntConstant(m, "DEF_MEM_LEVEL", DEF_MEM_LEVEL);
    PyModule_AddIntConstant(m, "Z_BEST_SPEED", Z_BEST_SPEED);
    PyModule_AddIntConstant(m, "Z_BEST_COMPRESSION", Z_BEST_COMPRESSION);
    PyModule_AddIntConstant(m, "Z_DEFAULT_COMPRESSION", Z_DEFAULT_COMPRESSION);
    PyModule_AddIntConstant(m, "Z_FILTERED", Z_FILTERED);
    PyModule_AddIntConstant(m, "Z_HUFFMAN_ONLY", Z_HUFFMAN_ONLY);
    PyModule_AddIntConstant(m, "Z_DEFAULT_STRATEGY", Z_DEFAULT_STRATEGY);
870

Jeremy Hylton's avatar
Jeremy Hylton committed
871 872 873 874
    PyModule_AddIntConstant(m, "Z_FINISH", Z_FINISH);
    PyModule_AddIntConstant(m, "Z_NO_FLUSH", Z_NO_FLUSH);
    PyModule_AddIntConstant(m, "Z_SYNC_FLUSH", Z_SYNC_FLUSH);
    PyModule_AddIntConstant(m, "Z_FULL_FLUSH", Z_FULL_FLUSH);
875

Jeremy Hylton's avatar
Jeremy Hylton committed
876 877 878 879 880
    ver = PyString_FromString(ZLIB_VERSION);
    if (ver != NULL) {
	PyDict_SetItemString(d, "ZLIB_VERSION", ver);
	Py_DECREF(ver);
    }
881 882

#ifdef WITH_THREAD
Jeremy Hylton's avatar
Jeremy Hylton committed
883
    zlib_lock = PyThread_allocate_lock();
884
#endif /* WITH_THREAD */
885
}