blob: 7c7c0e06517b095bde08002f71f264b2f5c5cd1f [file] [log] [blame]
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001/*
2
3python-bz2 - python bz2 library interface
4
5Copyright (c) 2002 Gustavo Niemeyer <niemeyer@conectiva.com>
6Copyright (c) 2002 Python Software Foundation; All Rights Reserved
7
8*/
9
Martin v. Löwise17af7b2002-11-23 09:16:19 +000010#include "Python.h"
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000011#include <stdio.h>
12#include <bzlib.h>
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000013#include "structmember.h"
14
15#ifdef WITH_THREAD
16#include "pythread.h"
17#endif
18
19static char __author__[] =
20"The bz2 python module was written by:\n\
21\n\
22 Gustavo Niemeyer <niemeyer@conectiva.com>\n\
23";
24
Georg Brandl33a5f2a2005-08-21 14:16:04 +000025/* Our very own off_t-like type, 64-bit if possible */
26/* copied from Objects/fileobject.c */
27#if !defined(HAVE_LARGEFILE_SUPPORT)
28typedef off_t Py_off_t;
29#elif SIZEOF_OFF_T >= 8
30typedef off_t Py_off_t;
31#elif SIZEOF_FPOS_T >= 8
32typedef fpos_t Py_off_t;
33#else
34#error "Large file support, but neither off_t nor fpos_t is large enough."
35#endif
36
Christian Heimes72b710a2008-05-26 13:28:38 +000037#define BUF(v) PyBytes_AS_STRING(v)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000038
39#define MODE_CLOSED 0
40#define MODE_READ 1
41#define MODE_READ_EOF 2
42#define MODE_WRITE 3
43
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000044#define BZ2FileObject_Check(v) (Py_TYPE(v) == &BZ2File_Type)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000045
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +000046
47#ifdef BZ_CONFIG_ERROR
48
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000049#if SIZEOF_LONG >= 8
50#define BZS_TOTAL_OUT(bzs) \
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000051 (((long)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000052#elif SIZEOF_LONG_LONG >= 8
53#define BZS_TOTAL_OUT(bzs) \
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000054 (((PY_LONG_LONG)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000055#else
56#define BZS_TOTAL_OUT(bzs) \
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000057 bzs->total_out_lo32
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000058#endif
59
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +000060#else /* ! BZ_CONFIG_ERROR */
61
62#define BZ2_bzRead bzRead
63#define BZ2_bzReadOpen bzReadOpen
64#define BZ2_bzReadClose bzReadClose
65#define BZ2_bzWrite bzWrite
66#define BZ2_bzWriteOpen bzWriteOpen
67#define BZ2_bzWriteClose bzWriteClose
68#define BZ2_bzCompress bzCompress
69#define BZ2_bzCompressInit bzCompressInit
70#define BZ2_bzCompressEnd bzCompressEnd
71#define BZ2_bzDecompress bzDecompress
72#define BZ2_bzDecompressInit bzDecompressInit
73#define BZ2_bzDecompressEnd bzDecompressEnd
74
75#define BZS_TOTAL_OUT(bzs) bzs->total_out
76
77#endif /* ! BZ_CONFIG_ERROR */
78
79
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000080#ifdef WITH_THREAD
Antoine Pitroua3af0eb2009-10-27 17:49:21 +000081#define ACQUIRE_LOCK(obj) do { \
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000082 if (!PyThread_acquire_lock(obj->lock, 0)) { \
83 Py_BEGIN_ALLOW_THREADS \
84 PyThread_acquire_lock(obj->lock, 1); \
85 Py_END_ALLOW_THREADS \
86 } } while(0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000087#define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)
88#else
89#define ACQUIRE_LOCK(obj)
90#define RELEASE_LOCK(obj)
91#endif
92
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000093/* Bits in f_newlinetypes */
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +000094#define NEWLINE_UNKNOWN 0 /* No newline seen, yet */
95#define NEWLINE_CR 1 /* \r newline seen */
96#define NEWLINE_LF 2 /* \n newline seen */
97#define NEWLINE_CRLF 4 /* \r\n newline seen */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000098
99/* ===================================================================== */
100/* Structure definitions. */
101
102typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000103 PyObject_HEAD
104 FILE *rawfp;
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000105
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000106 char* f_buf; /* Allocated readahead buffer */
107 char* f_bufend; /* Points after last occupied position */
108 char* f_bufptr; /* Current buffer position */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000109
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000110 BZFILE *fp;
111 int mode;
112 Py_off_t pos;
113 Py_off_t size;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000114#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000115 PyThread_type_lock lock;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000116#endif
117} BZ2FileObject;
118
119typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000120 PyObject_HEAD
121 bz_stream bzs;
122 int running;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000123#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000124 PyThread_type_lock lock;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000125#endif
126} BZ2CompObject;
127
128typedef struct {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000129 PyObject_HEAD
130 bz_stream bzs;
131 int running;
132 PyObject *unused_data;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000133#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000134 PyThread_type_lock lock;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000135#endif
136} BZ2DecompObject;
137
138/* ===================================================================== */
139/* Utility functions. */
140
141static int
142Util_CatchBZ2Error(int bzerror)
143{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000144 int ret = 0;
145 switch(bzerror) {
146 case BZ_OK:
147 case BZ_STREAM_END:
148 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000149
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +0000150#ifdef BZ_CONFIG_ERROR
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000151 case BZ_CONFIG_ERROR:
152 PyErr_SetString(PyExc_SystemError,
153 "the bz2 library was not compiled "
154 "correctly");
155 ret = 1;
156 break;
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +0000157#endif
Tim Peterse3228092002-11-09 04:21:44 +0000158
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000159 case BZ_PARAM_ERROR:
160 PyErr_SetString(PyExc_ValueError,
161 "the bz2 library has received wrong "
162 "parameters");
163 ret = 1;
164 break;
Tim Peterse3228092002-11-09 04:21:44 +0000165
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000166 case BZ_MEM_ERROR:
167 PyErr_NoMemory();
168 ret = 1;
169 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000170
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000171 case BZ_DATA_ERROR:
172 case BZ_DATA_ERROR_MAGIC:
173 PyErr_SetString(PyExc_IOError, "invalid data stream");
174 ret = 1;
175 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000176
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000177 case BZ_IO_ERROR:
178 PyErr_SetString(PyExc_IOError, "unknown IO error");
179 ret = 1;
180 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000181
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000182 case BZ_UNEXPECTED_EOF:
183 PyErr_SetString(PyExc_EOFError,
184 "compressed file ended before the "
185 "logical end-of-stream was detected");
186 ret = 1;
187 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000188
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000189 case BZ_SEQUENCE_ERROR:
190 PyErr_SetString(PyExc_RuntimeError,
191 "wrong sequence of bz2 library "
192 "commands used");
193 ret = 1;
194 break;
195 }
196 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000197}
198
199#if BUFSIZ < 8192
200#define SMALLCHUNK 8192
201#else
202#define SMALLCHUNK BUFSIZ
203#endif
204
205#if SIZEOF_INT < 4
206#define BIGCHUNK (512 * 32)
207#else
208#define BIGCHUNK (512 * 1024)
209#endif
210
211/* This is a hacked version of Python's fileobject.c:new_buffersize(). */
212static size_t
213Util_NewBufferSize(size_t currentsize)
214{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000215 if (currentsize > SMALLCHUNK) {
216 /* Keep doubling until we reach BIGCHUNK;
217 then keep adding BIGCHUNK. */
218 if (currentsize <= BIGCHUNK)
219 return currentsize + currentsize;
220 else
221 return currentsize + BIGCHUNK;
222 }
223 return currentsize + SMALLCHUNK;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000224}
225
226/* This is a hacked version of Python's fileobject.c:get_line(). */
227static PyObject *
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000228Util_GetLine(BZ2FileObject *f, int n)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000229{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000230 char c;
231 char *buf, *end;
232 size_t total_v_size; /* total # of slots in buffer */
233 size_t used_v_size; /* # used slots in buffer */
234 size_t increment; /* amount to increment the buffer */
235 PyObject *v;
236 int bzerror;
237 int bytes_read;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000238
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000239 total_v_size = n > 0 ? n : 100;
240 v = PyBytes_FromStringAndSize((char *)NULL, total_v_size);
241 if (v == NULL)
242 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000243
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000244 buf = BUF(v);
245 end = buf + total_v_size;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000246
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000247 for (;;) {
248 Py_BEGIN_ALLOW_THREADS
249 do {
250 bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
251 f->pos++;
252 if (bytes_read == 0)
253 break;
254 *buf++ = c;
255 } while (bzerror == BZ_OK && c != '\n' && buf != end);
256 Py_END_ALLOW_THREADS
257 if (bzerror == BZ_STREAM_END) {
258 f->size = f->pos;
259 f->mode = MODE_READ_EOF;
260 break;
261 } else if (bzerror != BZ_OK) {
262 Util_CatchBZ2Error(bzerror);
263 Py_DECREF(v);
264 return NULL;
265 }
266 if (c == '\n')
267 break;
268 /* Must be because buf == end */
269 if (n > 0)
270 break;
271 used_v_size = total_v_size;
272 increment = total_v_size >> 2; /* mild exponential growth */
273 total_v_size += increment;
274 if (total_v_size > INT_MAX) {
275 PyErr_SetString(PyExc_OverflowError,
276 "line is longer than a Python string can hold");
277 Py_DECREF(v);
278 return NULL;
279 }
280 if (_PyBytes_Resize(&v, total_v_size) < 0) {
281 return NULL;
282 }
283 buf = BUF(v) + used_v_size;
284 end = BUF(v) + total_v_size;
285 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000286
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000287 used_v_size = buf - BUF(v);
288 if (used_v_size != total_v_size) {
289 if (_PyBytes_Resize(&v, used_v_size) < 0) {
290 v = NULL;
291 }
292 }
293 return v;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000294}
295
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000296/* This is a hacked version of Python's fileobject.c:drop_readahead(). */
297static void
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000298Util_DropReadAhead(BZ2FileObject *f)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000299{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000300 if (f->f_buf != NULL) {
301 PyMem_Free(f->f_buf);
302 f->f_buf = NULL;
303 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000304}
305
306/* This is a hacked version of Python's fileobject.c:readahead(). */
307static int
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000308Util_ReadAhead(BZ2FileObject *f, int bufsize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000309{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000310 int chunksize;
311 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000312
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000313 if (f->f_buf != NULL) {
314 if((f->f_bufend - f->f_bufptr) >= 1)
315 return 0;
316 else
317 Util_DropReadAhead(f);
318 }
319 if (f->mode == MODE_READ_EOF) {
320 f->f_bufptr = f->f_buf;
321 f->f_bufend = f->f_buf;
322 return 0;
323 }
324 if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {
325 PyErr_NoMemory();
326 return -1;
327 }
328 Py_BEGIN_ALLOW_THREADS
329 chunksize = BZ2_bzRead(&bzerror, f->fp, f->f_buf, bufsize);
330 Py_END_ALLOW_THREADS
331 f->pos += chunksize;
332 if (bzerror == BZ_STREAM_END) {
333 f->size = f->pos;
334 f->mode = MODE_READ_EOF;
335 } else if (bzerror != BZ_OK) {
336 Util_CatchBZ2Error(bzerror);
337 Util_DropReadAhead(f);
338 return -1;
339 }
340 f->f_bufptr = f->f_buf;
341 f->f_bufend = f->f_buf + chunksize;
342 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000343}
344
345/* This is a hacked version of Python's
346 * fileobject.c:readahead_get_line_skip(). */
Christian Heimes72b710a2008-05-26 13:28:38 +0000347static PyBytesObject *
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000348Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000349{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000350 PyBytesObject* s;
351 char *bufptr;
352 char *buf;
353 int len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000354
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000355 if (f->f_buf == NULL)
356 if (Util_ReadAhead(f, bufsize) < 0)
357 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000358
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000359 len = f->f_bufend - f->f_bufptr;
360 if (len == 0)
361 return (PyBytesObject *)
362 PyBytes_FromStringAndSize(NULL, skip);
363 bufptr = memchr(f->f_bufptr, '\n', len);
364 if (bufptr != NULL) {
365 bufptr++; /* Count the '\n' */
366 len = bufptr - f->f_bufptr;
367 s = (PyBytesObject *)
368 PyBytes_FromStringAndSize(NULL, skip+len);
369 if (s == NULL)
370 return NULL;
371 memcpy(PyBytes_AS_STRING(s)+skip, f->f_bufptr, len);
372 f->f_bufptr = bufptr;
373 if (bufptr == f->f_bufend)
374 Util_DropReadAhead(f);
375 } else {
376 bufptr = f->f_bufptr;
377 buf = f->f_buf;
378 f->f_buf = NULL; /* Force new readahead buffer */
379 s = Util_ReadAheadGetLineSkip(f, skip+len,
380 bufsize + (bufsize>>2));
381 if (s == NULL) {
382 PyMem_Free(buf);
383 return NULL;
384 }
385 memcpy(PyBytes_AS_STRING(s)+skip, bufptr, len);
386 PyMem_Free(buf);
387 }
388 return s;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000389}
390
391/* ===================================================================== */
392/* Methods of BZ2File. */
393
394PyDoc_STRVAR(BZ2File_read__doc__,
395"read([size]) -> string\n\
396\n\
397Read at most size uncompressed bytes, returned as a string. If the size\n\
398argument is negative or omitted, read until EOF is reached.\n\
399");
400
401/* This is a hacked version of Python's fileobject.c:file_read(). */
402static PyObject *
403BZ2File_read(BZ2FileObject *self, PyObject *args)
404{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000405 long bytesrequested = -1;
406 size_t bytesread, buffersize, chunksize;
407 int bzerror;
408 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000409
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000410 if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))
411 return NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000412
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000413 ACQUIRE_LOCK(self);
414 switch (self->mode) {
415 case MODE_READ:
416 break;
417 case MODE_READ_EOF:
418 ret = PyBytes_FromStringAndSize("", 0);
419 goto cleanup;
420 case MODE_CLOSED:
421 PyErr_SetString(PyExc_ValueError,
422 "I/O operation on closed file");
423 goto cleanup;
424 default:
425 PyErr_SetString(PyExc_IOError,
426 "file is not ready for reading");
427 goto cleanup;
428 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000429
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000430 if (bytesrequested < 0)
431 buffersize = Util_NewBufferSize((size_t)0);
432 else
433 buffersize = bytesrequested;
434 if (buffersize > INT_MAX) {
435 PyErr_SetString(PyExc_OverflowError,
436 "requested number of bytes is "
437 "more than a Python string can hold");
438 goto cleanup;
439 }
440 ret = PyBytes_FromStringAndSize((char *)NULL, buffersize);
441 if (ret == NULL || buffersize == 0)
442 goto cleanup;
443 bytesread = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000444
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000445 for (;;) {
446 Py_BEGIN_ALLOW_THREADS
447 chunksize = BZ2_bzRead(&bzerror, self->fp,
448 BUF(ret)+bytesread,
449 buffersize-bytesread);
450 self->pos += chunksize;
451 Py_END_ALLOW_THREADS
452 bytesread += chunksize;
453 if (bzerror == BZ_STREAM_END) {
454 self->size = self->pos;
455 self->mode = MODE_READ_EOF;
456 break;
457 } else if (bzerror != BZ_OK) {
458 Util_CatchBZ2Error(bzerror);
459 Py_DECREF(ret);
460 ret = NULL;
461 goto cleanup;
462 }
463 if (bytesrequested < 0) {
464 buffersize = Util_NewBufferSize(buffersize);
465 if (_PyBytes_Resize(&ret, buffersize) < 0) {
466 ret = NULL;
467 goto cleanup;
468 }
469 } else {
470 break;
471 }
472 }
473 if (bytesread != buffersize) {
474 if (_PyBytes_Resize(&ret, bytesread) < 0) {
475 ret = NULL;
476 }
477 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000478
479cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000480 RELEASE_LOCK(self);
481 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000482}
483
484PyDoc_STRVAR(BZ2File_readline__doc__,
485"readline([size]) -> string\n\
486\n\
487Return the next line from the file, as a string, retaining newline.\n\
488A non-negative size argument will limit the maximum number of bytes to\n\
489return (an incomplete line may be returned then). Return an empty\n\
490string at EOF.\n\
491");
492
493static PyObject *
494BZ2File_readline(BZ2FileObject *self, PyObject *args)
495{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000496 PyObject *ret = NULL;
497 int sizehint = -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000498
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000499 if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))
500 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000501
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000502 ACQUIRE_LOCK(self);
503 switch (self->mode) {
504 case MODE_READ:
505 break;
506 case MODE_READ_EOF:
507 ret = PyBytes_FromStringAndSize("", 0);
508 goto cleanup;
509 case MODE_CLOSED:
510 PyErr_SetString(PyExc_ValueError,
511 "I/O operation on closed file");
512 goto cleanup;
513 default:
514 PyErr_SetString(PyExc_IOError,
515 "file is not ready for reading");
516 goto cleanup;
517 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000518
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000519 if (sizehint == 0)
520 ret = PyBytes_FromStringAndSize("", 0);
521 else
522 ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000523
524cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000525 RELEASE_LOCK(self);
526 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000527}
528
529PyDoc_STRVAR(BZ2File_readlines__doc__,
530"readlines([size]) -> list\n\
531\n\
532Call readline() repeatedly and return a list of lines read.\n\
533The optional size argument, if given, is an approximate bound on the\n\
534total number of bytes in the lines returned.\n\
535");
536
537/* This is a hacked version of Python's fileobject.c:file_readlines(). */
538static PyObject *
539BZ2File_readlines(BZ2FileObject *self, PyObject *args)
540{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000541 long sizehint = 0;
542 PyObject *list = NULL;
543 PyObject *line;
544 char small_buffer[SMALLCHUNK];
545 char *buffer = small_buffer;
546 size_t buffersize = SMALLCHUNK;
547 PyObject *big_buffer = NULL;
548 size_t nfilled = 0;
549 size_t nread;
550 size_t totalread = 0;
551 char *p, *q, *end;
552 int err;
553 int shortread = 0;
554 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000555
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000556 if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))
557 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000558
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000559 ACQUIRE_LOCK(self);
560 switch (self->mode) {
561 case MODE_READ:
562 break;
563 case MODE_READ_EOF:
564 list = PyList_New(0);
565 goto cleanup;
566 case MODE_CLOSED:
567 PyErr_SetString(PyExc_ValueError,
568 "I/O operation on closed file");
569 goto cleanup;
570 default:
571 PyErr_SetString(PyExc_IOError,
572 "file is not ready for reading");
573 goto cleanup;
574 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000575
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000576 if ((list = PyList_New(0)) == NULL)
577 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000578
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000579 for (;;) {
580 Py_BEGIN_ALLOW_THREADS
581 nread = BZ2_bzRead(&bzerror, self->fp,
582 buffer+nfilled, buffersize-nfilled);
583 self->pos += nread;
584 Py_END_ALLOW_THREADS
585 if (bzerror == BZ_STREAM_END) {
586 self->size = self->pos;
587 self->mode = MODE_READ_EOF;
588 if (nread == 0) {
589 sizehint = 0;
590 break;
591 }
592 shortread = 1;
593 } else if (bzerror != BZ_OK) {
594 Util_CatchBZ2Error(bzerror);
595 error:
596 Py_DECREF(list);
597 list = NULL;
598 goto cleanup;
599 }
600 totalread += nread;
601 p = memchr(buffer+nfilled, '\n', nread);
602 if (!shortread && p == NULL) {
603 /* Need a larger buffer to fit this line */
604 nfilled += nread;
605 buffersize *= 2;
606 if (buffersize > INT_MAX) {
607 PyErr_SetString(PyExc_OverflowError,
608 "line is longer than a Python string can hold");
609 goto error;
610 }
611 if (big_buffer == NULL) {
612 /* Create the big buffer */
613 big_buffer = PyBytes_FromStringAndSize(
614 NULL, buffersize);
615 if (big_buffer == NULL)
616 goto error;
617 buffer = PyBytes_AS_STRING(big_buffer);
618 memcpy(buffer, small_buffer, nfilled);
619 }
620 else {
621 /* Grow the big buffer */
622 if (_PyBytes_Resize(&big_buffer, buffersize) < 0){
623 big_buffer = NULL;
624 goto error;
625 }
626 buffer = PyBytes_AS_STRING(big_buffer);
627 }
628 continue;
629 }
630 end = buffer+nfilled+nread;
631 q = buffer;
632 while (p != NULL) {
633 /* Process complete lines */
634 p++;
635 line = PyBytes_FromStringAndSize(q, p-q);
636 if (line == NULL)
637 goto error;
638 err = PyList_Append(list, line);
639 Py_DECREF(line);
640 if (err != 0)
641 goto error;
642 q = p;
643 p = memchr(q, '\n', end-q);
644 }
645 /* Move the remaining incomplete line to the start */
646 nfilled = end-q;
647 memmove(buffer, q, nfilled);
648 if (sizehint > 0)
649 if (totalread >= (size_t)sizehint)
650 break;
651 if (shortread) {
652 sizehint = 0;
653 break;
654 }
655 }
656 if (nfilled != 0) {
657 /* Partial last line */
658 line = PyBytes_FromStringAndSize(buffer, nfilled);
659 if (line == NULL)
660 goto error;
661 if (sizehint > 0) {
662 /* Need to complete the last line */
663 PyObject *rest = Util_GetLine(self, 0);
664 if (rest == NULL) {
665 Py_DECREF(line);
666 goto error;
667 }
668 PyBytes_Concat(&line, rest);
669 Py_DECREF(rest);
670 if (line == NULL)
671 goto error;
672 }
673 err = PyList_Append(list, line);
674 Py_DECREF(line);
675 if (err != 0)
676 goto error;
677 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000678
679 cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000680 RELEASE_LOCK(self);
681 if (big_buffer) {
682 Py_DECREF(big_buffer);
683 }
684 return list;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000685}
686
687PyDoc_STRVAR(BZ2File_write__doc__,
688"write(data) -> None\n\
689\n\
690Write the 'data' string to file. Note that due to buffering, close() may\n\
691be needed before the file on disk reflects the data written.\n\
692");
693
694/* This is a hacked version of Python's fileobject.c:file_write(). */
695static PyObject *
696BZ2File_write(BZ2FileObject *self, PyObject *args)
697{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000698 PyObject *ret = NULL;
699 Py_buffer pbuf;
700 char *buf;
701 int len;
702 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000703
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000704 if (!PyArg_ParseTuple(args, "y*:write", &pbuf))
705 return NULL;
706 buf = pbuf.buf;
707 len = pbuf.len;
Tim Peterse3228092002-11-09 04:21:44 +0000708
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000709 ACQUIRE_LOCK(self);
710 switch (self->mode) {
711 case MODE_WRITE:
712 break;
Tim Peterse3228092002-11-09 04:21:44 +0000713
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000714 case MODE_CLOSED:
715 PyErr_SetString(PyExc_ValueError,
716 "I/O operation on closed file");
717 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000718
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000719 default:
720 PyErr_SetString(PyExc_IOError,
721 "file is not ready for writing");
722 goto cleanup;
723 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000724
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000725 Py_BEGIN_ALLOW_THREADS
726 BZ2_bzWrite (&bzerror, self->fp, buf, len);
727 self->pos += len;
728 Py_END_ALLOW_THREADS
Tim Peterse3228092002-11-09 04:21:44 +0000729
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000730 if (bzerror != BZ_OK) {
731 Util_CatchBZ2Error(bzerror);
732 goto cleanup;
733 }
Tim Peterse3228092002-11-09 04:21:44 +0000734
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000735 Py_INCREF(Py_None);
736 ret = Py_None;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000737
738cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000739 PyBuffer_Release(&pbuf);
740 RELEASE_LOCK(self);
741 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000742}
743
744PyDoc_STRVAR(BZ2File_writelines__doc__,
745"writelines(sequence_of_strings) -> None\n\
746\n\
747Write the sequence of strings to the file. Note that newlines are not\n\
748added. The sequence can be any iterable object producing strings. This is\n\
749equivalent to calling write() for each string.\n\
750");
751
752/* This is a hacked version of Python's fileobject.c:file_writelines(). */
753static PyObject *
754BZ2File_writelines(BZ2FileObject *self, PyObject *seq)
755{
756#define CHUNKSIZE 1000
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000757 PyObject *list = NULL;
758 PyObject *iter = NULL;
759 PyObject *ret = NULL;
760 PyObject *line;
761 int i, j, index, len, islist;
762 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000763
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000764 ACQUIRE_LOCK(self);
765 switch (self->mode) {
766 case MODE_WRITE:
767 break;
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000768
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000769 case MODE_CLOSED:
770 PyErr_SetString(PyExc_ValueError,
771 "I/O operation on closed file");
772 goto error;
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000773
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000774 default:
775 PyErr_SetString(PyExc_IOError,
776 "file is not ready for writing");
777 goto error;
778 }
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000779
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000780 islist = PyList_Check(seq);
781 if (!islist) {
782 iter = PyObject_GetIter(seq);
783 if (iter == NULL) {
784 PyErr_SetString(PyExc_TypeError,
785 "writelines() requires an iterable argument");
786 goto error;
787 }
788 list = PyList_New(CHUNKSIZE);
789 if (list == NULL)
790 goto error;
791 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000792
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000793 /* Strategy: slurp CHUNKSIZE lines into a private list,
794 checking that they are all strings, then write that list
795 without holding the interpreter lock, then come back for more. */
796 for (index = 0; ; index += CHUNKSIZE) {
797 if (islist) {
798 Py_XDECREF(list);
799 list = PyList_GetSlice(seq, index, index+CHUNKSIZE);
800 if (list == NULL)
801 goto error;
802 j = PyList_GET_SIZE(list);
803 }
804 else {
805 for (j = 0; j < CHUNKSIZE; j++) {
806 line = PyIter_Next(iter);
807 if (line == NULL) {
808 if (PyErr_Occurred())
809 goto error;
810 break;
811 }
812 PyList_SetItem(list, j, line);
813 }
814 }
815 if (j == 0)
816 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000817
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000818 /* Check that all entries are indeed byte strings. If not,
819 apply the same rules as for file.write() and
820 convert the rets to strings. This is slow, but
821 seems to be the only way since all conversion APIs
822 could potentially execute Python code. */
823 for (i = 0; i < j; i++) {
824 PyObject *v = PyList_GET_ITEM(list, i);
825 if (!PyBytes_Check(v)) {
826 const char *buffer;
827 Py_ssize_t len;
828 if (PyObject_AsCharBuffer(v, &buffer, &len)) {
829 PyErr_SetString(PyExc_TypeError,
830 "writelines() "
831 "argument must be "
832 "a sequence of "
833 "bytes objects");
834 goto error;
835 }
836 line = PyBytes_FromStringAndSize(buffer,
837 len);
838 if (line == NULL)
839 goto error;
840 Py_DECREF(v);
841 PyList_SET_ITEM(list, i, line);
842 }
843 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000844
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000845 /* Since we are releasing the global lock, the
846 following code may *not* execute Python code. */
847 Py_BEGIN_ALLOW_THREADS
848 for (i = 0; i < j; i++) {
849 line = PyList_GET_ITEM(list, i);
850 len = PyBytes_GET_SIZE(line);
851 BZ2_bzWrite (&bzerror, self->fp,
852 PyBytes_AS_STRING(line), len);
853 if (bzerror != BZ_OK) {
854 Py_BLOCK_THREADS
855 Util_CatchBZ2Error(bzerror);
856 goto error;
857 }
858 }
859 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000860
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000861 if (j < CHUNKSIZE)
862 break;
863 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000864
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000865 Py_INCREF(Py_None);
866 ret = Py_None;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000867
868 error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000869 RELEASE_LOCK(self);
870 Py_XDECREF(list);
871 Py_XDECREF(iter);
872 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000873#undef CHUNKSIZE
874}
875
876PyDoc_STRVAR(BZ2File_seek__doc__,
877"seek(offset [, whence]) -> None\n\
878\n\
879Move to new file position. Argument offset is a byte count. Optional\n\
880argument whence defaults to 0 (offset from start of file, offset\n\
881should be >= 0); other values are 1 (move relative to current position,\n\
882positive or negative), and 2 (move relative to end of file, usually\n\
883negative, although many platforms allow seeking beyond the end of a file).\n\
884\n\
885Note that seeking of bz2 files is emulated, and depending on the parameters\n\
886the operation may be extremely slow.\n\
887");
888
889static PyObject *
890BZ2File_seek(BZ2FileObject *self, PyObject *args)
891{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000892 int where = 0;
893 PyObject *offobj;
894 Py_off_t offset;
895 char small_buffer[SMALLCHUNK];
896 char *buffer = small_buffer;
897 size_t buffersize = SMALLCHUNK;
898 Py_off_t bytesread = 0;
899 size_t readsize;
900 int chunksize;
901 int bzerror;
902 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000903
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000904 if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))
905 return NULL;
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000906#if !defined(HAVE_LARGEFILE_SUPPORT)
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000907 offset = PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000908#else
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000909 offset = PyLong_Check(offobj) ?
910 PyLong_AsLongLong(offobj) : PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000911#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000912 if (PyErr_Occurred())
913 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000914
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000915 ACQUIRE_LOCK(self);
916 Util_DropReadAhead(self);
917 switch (self->mode) {
918 case MODE_READ:
919 case MODE_READ_EOF:
920 break;
Tim Peterse3228092002-11-09 04:21:44 +0000921
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000922 case MODE_CLOSED:
923 PyErr_SetString(PyExc_ValueError,
924 "I/O operation on closed file");
925 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000926
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000927 default:
928 PyErr_SetString(PyExc_IOError,
929 "seek works only while reading");
930 goto cleanup;
931 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000932
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000933 if (where == 2) {
934 if (self->size == -1) {
935 assert(self->mode != MODE_READ_EOF);
936 for (;;) {
937 Py_BEGIN_ALLOW_THREADS
938 chunksize = BZ2_bzRead(&bzerror, self->fp,
939 buffer, buffersize);
940 self->pos += chunksize;
941 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000942
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000943 bytesread += chunksize;
944 if (bzerror == BZ_STREAM_END) {
945 break;
946 } else if (bzerror != BZ_OK) {
947 Util_CatchBZ2Error(bzerror);
948 goto cleanup;
949 }
950 }
951 self->mode = MODE_READ_EOF;
952 self->size = self->pos;
953 bytesread = 0;
954 }
955 offset = self->size + offset;
956 } else if (where == 1) {
957 offset = self->pos + offset;
958 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000959
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000960 /* Before getting here, offset must be the absolute position the file
961 * pointer should be set to. */
Georg Brandl47fab922006-02-18 21:57:25 +0000962
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000963 if (offset >= self->pos) {
964 /* we can move forward */
965 offset -= self->pos;
966 } else {
967 /* we cannot move back, so rewind the stream */
968 BZ2_bzReadClose(&bzerror, self->fp);
969 if (bzerror != BZ_OK) {
970 Util_CatchBZ2Error(bzerror);
971 goto cleanup;
972 }
973 rewind(self->rawfp);
974 self->pos = 0;
975 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
976 0, 0, NULL, 0);
977 if (bzerror != BZ_OK) {
978 Util_CatchBZ2Error(bzerror);
979 goto cleanup;
980 }
981 self->mode = MODE_READ;
982 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000983
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000984 if (offset <= 0 || self->mode == MODE_READ_EOF)
985 goto exit;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000986
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +0000987 /* Before getting here, offset must be set to the number of bytes
988 * to walk forward. */
989 for (;;) {
990 if (offset-bytesread > buffersize)
991 readsize = buffersize;
992 else
993 /* offset might be wider that readsize, but the result
994 * of the subtraction is bound by buffersize (see the
995 * condition above). buffersize is 8192. */
996 readsize = (size_t)(offset-bytesread);
997 Py_BEGIN_ALLOW_THREADS
998 chunksize = BZ2_bzRead(&bzerror, self->fp, buffer, readsize);
999 self->pos += chunksize;
1000 Py_END_ALLOW_THREADS
1001 bytesread += chunksize;
1002 if (bzerror == BZ_STREAM_END) {
1003 self->size = self->pos;
1004 self->mode = MODE_READ_EOF;
1005 break;
1006 } else if (bzerror != BZ_OK) {
1007 Util_CatchBZ2Error(bzerror);
1008 goto cleanup;
1009 }
1010 if (bytesread == offset)
1011 break;
1012 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001013
1014exit:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001015 Py_INCREF(Py_None);
1016 ret = Py_None;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001017
1018cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001019 RELEASE_LOCK(self);
1020 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001021}
1022
1023PyDoc_STRVAR(BZ2File_tell__doc__,
1024"tell() -> int\n\
1025\n\
1026Return the current file position, an integer (may be a long integer).\n\
1027");
1028
1029static PyObject *
1030BZ2File_tell(BZ2FileObject *self, PyObject *args)
1031{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001032 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001033
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001034 if (self->mode == MODE_CLOSED) {
1035 PyErr_SetString(PyExc_ValueError,
1036 "I/O operation on closed file");
1037 goto cleanup;
1038 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001039
Georg Brandla8bcecc2005-09-03 07:49:53 +00001040#if !defined(HAVE_LARGEFILE_SUPPORT)
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001041 ret = PyLong_FromLong(self->pos);
Georg Brandla8bcecc2005-09-03 07:49:53 +00001042#else
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001043 ret = PyLong_FromLongLong(self->pos);
Georg Brandla8bcecc2005-09-03 07:49:53 +00001044#endif
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001045
1046cleanup:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001047 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001048}
1049
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001050PyDoc_STRVAR(BZ2File_close__doc__,
1051"close() -> None or (perhaps) an integer\n\
1052\n\
1053Close the file. Sets data attribute .closed to true. A closed file\n\
1054cannot be used for further I/O operations. close() may be called more\n\
1055than once without error.\n\
1056");
1057
1058static PyObject *
1059BZ2File_close(BZ2FileObject *self)
1060{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001061 PyObject *ret = NULL;
1062 int bzerror = BZ_OK;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001063
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001064 if (self->mode == MODE_CLOSED) {
1065 Py_RETURN_NONE;
1066 }
Guido van Rossumf09ca142007-06-13 00:03:05 +00001067
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001068 ACQUIRE_LOCK(self);
1069 switch (self->mode) {
1070 case MODE_READ:
1071 case MODE_READ_EOF:
1072 BZ2_bzReadClose(&bzerror, self->fp);
1073 break;
1074 case MODE_WRITE:
1075 BZ2_bzWriteClose(&bzerror, self->fp,
1076 0, NULL, NULL);
1077 break;
1078 }
1079 self->mode = MODE_CLOSED;
1080 fclose(self->rawfp);
1081 self->rawfp = NULL;
1082 if (bzerror == BZ_OK) {
1083 Py_INCREF(Py_None);
1084 ret = Py_None;
1085 }
1086 else {
1087 Util_CatchBZ2Error(bzerror);
1088 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001089
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001090 RELEASE_LOCK(self);
1091 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001092}
1093
Antoine Pitrou308705e2009-01-10 16:22:51 +00001094PyDoc_STRVAR(BZ2File_enter_doc,
1095"__enter__() -> self.");
1096
1097static PyObject *
1098BZ2File_enter(BZ2FileObject *self)
1099{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001100 if (self->mode == MODE_CLOSED) {
1101 PyErr_SetString(PyExc_ValueError,
1102 "I/O operation on closed file");
1103 return NULL;
1104 }
1105 Py_INCREF(self);
1106 return (PyObject *) self;
Antoine Pitrou308705e2009-01-10 16:22:51 +00001107}
1108
1109PyDoc_STRVAR(BZ2File_exit_doc,
1110"__exit__(*excinfo) -> None. Closes the file.");
1111
1112static PyObject *
1113BZ2File_exit(BZ2FileObject *self, PyObject *args)
1114{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001115 PyObject *ret = PyObject_CallMethod((PyObject *) self, "close", NULL);
1116 if (!ret)
1117 /* If error occurred, pass through */
1118 return NULL;
1119 Py_DECREF(ret);
1120 Py_RETURN_NONE;
Antoine Pitrou308705e2009-01-10 16:22:51 +00001121}
1122
1123
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001124static PyObject *BZ2File_getiter(BZ2FileObject *self);
1125
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001126static PyMethodDef BZ2File_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001127 {"read", (PyCFunction)BZ2File_read, METH_VARARGS, BZ2File_read__doc__},
1128 {"readline", (PyCFunction)BZ2File_readline, METH_VARARGS, BZ2File_readline__doc__},
1129 {"readlines", (PyCFunction)BZ2File_readlines, METH_VARARGS, BZ2File_readlines__doc__},
1130 {"write", (PyCFunction)BZ2File_write, METH_VARARGS, BZ2File_write__doc__},
1131 {"writelines", (PyCFunction)BZ2File_writelines, METH_O, BZ2File_writelines__doc__},
1132 {"seek", (PyCFunction)BZ2File_seek, METH_VARARGS, BZ2File_seek__doc__},
1133 {"tell", (PyCFunction)BZ2File_tell, METH_NOARGS, BZ2File_tell__doc__},
1134 {"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
1135 {"__enter__", (PyCFunction)BZ2File_enter, METH_NOARGS, BZ2File_enter_doc},
1136 {"__exit__", (PyCFunction)BZ2File_exit, METH_VARARGS, BZ2File_exit_doc},
1137 {NULL, NULL} /* sentinel */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001138};
1139
1140
1141/* ===================================================================== */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001142/* Getters and setters of BZ2File. */
1143
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001144static PyObject *
1145BZ2File_get_closed(BZ2FileObject *self, void *closure)
1146{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001147 return PyLong_FromLong(self->mode == MODE_CLOSED);
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001148}
1149
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001150static PyGetSetDef BZ2File_getset[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001151 {"closed", (getter)BZ2File_get_closed, NULL,
1152 "True if the file is closed"},
1153 {NULL} /* Sentinel */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001154};
1155
1156
1157/* ===================================================================== */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001158/* Slot definitions for BZ2File_Type. */
1159
1160static int
1161BZ2File_init(BZ2FileObject *self, PyObject *args, PyObject *kwargs)
1162{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001163 static char *kwlist[] = {"filename", "mode", "buffering",
1164 "compresslevel", 0};
1165 PyObject *name_obj = NULL;
1166 char *name;
1167 char *mode = "r";
1168 int buffering = -1;
1169 int compresslevel = 9;
1170 int bzerror;
1171 int mode_char = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001172
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001173 self->size = -1;
Tim Peterse3228092002-11-09 04:21:44 +00001174
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001175 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&|sii:BZ2File",
1176 kwlist, PyUnicode_FSConverter, &name_obj,
1177 &mode, &buffering,
1178 &compresslevel))
1179 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001180
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001181 if (PyBytes_Check(name_obj))
1182 name = PyBytes_AsString(name_obj);
1183 else
1184 name = PyByteArray_AsString(name_obj);
1185 if (compresslevel < 1 || compresslevel > 9) {
1186 PyErr_SetString(PyExc_ValueError,
1187 "compresslevel must be between 1 and 9");
1188 Py_DECREF(name_obj);
1189 return -1;
1190 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001191
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001192 for (;;) {
1193 int error = 0;
1194 switch (*mode) {
1195 case 'r':
1196 case 'w':
1197 if (mode_char)
1198 error = 1;
1199 mode_char = *mode;
1200 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001201
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001202 case 'b':
1203 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001204
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001205 default:
1206 error = 1;
1207 break;
1208 }
1209 if (error) {
1210 PyErr_Format(PyExc_ValueError,
1211 "invalid mode char %c", *mode);
1212 Py_DECREF(name_obj);
1213 return -1;
1214 }
1215 mode++;
1216 if (*mode == '\0')
1217 break;
1218 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001219
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001220 if (mode_char == 0) {
1221 mode_char = 'r';
1222 }
Georg Brandl6b95f1d2005-06-03 19:47:00 +00001223
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001224 mode = (mode_char == 'r') ? "rb" : "wb";
Tim Peterse3228092002-11-09 04:21:44 +00001225
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001226 self->rawfp = fopen(name, mode);
1227 Py_DECREF(name_obj);
1228 if (self->rawfp == NULL) {
1229 PyErr_SetFromErrno(PyExc_IOError);
1230 return -1;
1231 }
1232 /* XXX Ignore buffering */
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001233
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001234 /* From now on, we have stuff to dealloc, so jump to error label
1235 * instead of returning */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001236
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001237#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001238 self->lock = PyThread_allocate_lock();
1239 if (!self->lock) {
1240 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1241 goto error;
1242 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001243#endif
1244
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001245 if (mode_char == 'r')
1246 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
1247 0, 0, NULL, 0);
1248 else
1249 self->fp = BZ2_bzWriteOpen(&bzerror, self->rawfp,
1250 compresslevel, 0, 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001251
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001252 if (bzerror != BZ_OK) {
1253 Util_CatchBZ2Error(bzerror);
1254 goto error;
1255 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001256
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001257 self->mode = (mode_char == 'r') ? MODE_READ : MODE_WRITE;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001258
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001259 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001260
1261error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001262 fclose(self->rawfp);
1263 self->rawfp = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001264#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001265 if (self->lock) {
1266 PyThread_free_lock(self->lock);
1267 self->lock = NULL;
1268 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001269#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001270 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001271}
1272
1273static void
1274BZ2File_dealloc(BZ2FileObject *self)
1275{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001276 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001277#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001278 if (self->lock)
1279 PyThread_free_lock(self->lock);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001280#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001281 switch (self->mode) {
1282 case MODE_READ:
1283 case MODE_READ_EOF:
1284 BZ2_bzReadClose(&bzerror, self->fp);
1285 break;
1286 case MODE_WRITE:
1287 BZ2_bzWriteClose(&bzerror, self->fp,
1288 0, NULL, NULL);
1289 break;
1290 }
1291 Util_DropReadAhead(self);
1292 if (self->rawfp != NULL)
1293 fclose(self->rawfp);
1294 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001295}
1296
1297/* This is a hacked version of Python's fileobject.c:file_getiter(). */
1298static PyObject *
1299BZ2File_getiter(BZ2FileObject *self)
1300{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001301 if (self->mode == MODE_CLOSED) {
1302 PyErr_SetString(PyExc_ValueError,
1303 "I/O operation on closed file");
1304 return NULL;
1305 }
1306 Py_INCREF((PyObject*)self);
1307 return (PyObject *)self;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001308}
1309
1310/* This is a hacked version of Python's fileobject.c:file_iternext(). */
1311#define READAHEAD_BUFSIZE 8192
1312static PyObject *
1313BZ2File_iternext(BZ2FileObject *self)
1314{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001315 PyBytesObject* ret;
1316 ACQUIRE_LOCK(self);
1317 if (self->mode == MODE_CLOSED) {
1318 RELEASE_LOCK(self);
1319 PyErr_SetString(PyExc_ValueError,
1320 "I/O operation on closed file");
1321 return NULL;
1322 }
1323 ret = Util_ReadAheadGetLineSkip(self, 0, READAHEAD_BUFSIZE);
1324 RELEASE_LOCK(self);
1325 if (ret == NULL || PyBytes_GET_SIZE(ret) == 0) {
1326 Py_XDECREF(ret);
1327 return NULL;
1328 }
1329 return (PyObject *)ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001330}
1331
1332/* ===================================================================== */
1333/* BZ2File_Type definition. */
1334
1335PyDoc_VAR(BZ2File__doc__) =
1336PyDoc_STR(
1337"BZ2File(name [, mode='r', buffering=0, compresslevel=9]) -> file object\n\
1338\n\
1339Open a bz2 file. The mode can be 'r' or 'w', for reading (default) or\n\
1340writing. When opened for writing, the file will be created if it doesn't\n\
1341exist, and truncated otherwise. If the buffering argument is given, 0 means\n\
1342unbuffered, and larger numbers specify the buffer size. If compresslevel\n\
1343is given, must be a number between 1 and 9.\n\
Guido van Rossum88e860c2007-06-13 01:46:31 +00001344Data read is always returned in bytes; data written ought to be bytes.\n\
1345");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001346
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001347static PyTypeObject BZ2File_Type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001348 PyVarObject_HEAD_INIT(NULL, 0)
1349 "bz2.BZ2File", /*tp_name*/
1350 sizeof(BZ2FileObject), /*tp_basicsize*/
1351 0, /*tp_itemsize*/
1352 (destructor)BZ2File_dealloc, /*tp_dealloc*/
1353 0, /*tp_print*/
1354 0, /*tp_getattr*/
1355 0, /*tp_setattr*/
1356 0, /*tp_reserved*/
1357 0, /*tp_repr*/
1358 0, /*tp_as_number*/
1359 0, /*tp_as_sequence*/
1360 0, /*tp_as_mapping*/
1361 0, /*tp_hash*/
1362 0, /*tp_call*/
1363 0, /*tp_str*/
1364 PyObject_GenericGetAttr,/*tp_getattro*/
1365 PyObject_GenericSetAttr,/*tp_setattro*/
1366 0, /*tp_as_buffer*/
1367 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1368 BZ2File__doc__, /*tp_doc*/
1369 0, /*tp_traverse*/
1370 0, /*tp_clear*/
1371 0, /*tp_richcompare*/
1372 0, /*tp_weaklistoffset*/
1373 (getiterfunc)BZ2File_getiter, /*tp_iter*/
1374 (iternextfunc)BZ2File_iternext, /*tp_iternext*/
1375 BZ2File_methods, /*tp_methods*/
1376 0, /*tp_members*/
1377 BZ2File_getset, /*tp_getset*/
1378 0, /*tp_base*/
1379 0, /*tp_dict*/
1380 0, /*tp_descr_get*/
1381 0, /*tp_descr_set*/
1382 0, /*tp_dictoffset*/
1383 (initproc)BZ2File_init, /*tp_init*/
1384 PyType_GenericAlloc, /*tp_alloc*/
1385 PyType_GenericNew, /*tp_new*/
1386 PyObject_Free, /*tp_free*/
1387 0, /*tp_is_gc*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001388};
1389
1390
1391/* ===================================================================== */
1392/* Methods of BZ2Comp. */
1393
1394PyDoc_STRVAR(BZ2Comp_compress__doc__,
1395"compress(data) -> string\n\
1396\n\
1397Provide more data to the compressor object. It will return chunks of\n\
1398compressed data whenever possible. When you've finished providing data\n\
1399to compress, call the flush() method to finish the compression process,\n\
1400and return what is left in the internal buffers.\n\
1401");
1402
1403static PyObject *
1404BZ2Comp_compress(BZ2CompObject *self, PyObject *args)
1405{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001406 Py_buffer pdata;
1407 char *data;
1408 int datasize;
1409 int bufsize = SMALLCHUNK;
1410 PY_LONG_LONG totalout;
1411 PyObject *ret = NULL;
1412 bz_stream *bzs = &self->bzs;
1413 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001414
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001415 if (!PyArg_ParseTuple(args, "y*:compress", &pdata))
1416 return NULL;
1417 data = pdata.buf;
1418 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001419
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001420 if (datasize == 0) {
1421 PyBuffer_Release(&pdata);
1422 return PyBytes_FromStringAndSize("", 0);
1423 }
Gustavo Niemeyera6e436e2004-02-14 00:02:45 +00001424
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001425 ACQUIRE_LOCK(self);
1426 if (!self->running) {
1427 PyErr_SetString(PyExc_ValueError,
1428 "this object was already flushed");
1429 goto error;
1430 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001431
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001432 ret = PyBytes_FromStringAndSize(NULL, bufsize);
1433 if (!ret)
1434 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001435
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001436 bzs->next_in = data;
1437 bzs->avail_in = datasize;
1438 bzs->next_out = BUF(ret);
1439 bzs->avail_out = bufsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001440
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001441 totalout = BZS_TOTAL_OUT(bzs);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001442
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001443 for (;;) {
1444 Py_BEGIN_ALLOW_THREADS
1445 bzerror = BZ2_bzCompress(bzs, BZ_RUN);
1446 Py_END_ALLOW_THREADS
1447 if (bzerror != BZ_RUN_OK) {
1448 Util_CatchBZ2Error(bzerror);
1449 goto error;
1450 }
1451 if (bzs->avail_in == 0)
1452 break; /* no more input data */
1453 if (bzs->avail_out == 0) {
1454 bufsize = Util_NewBufferSize(bufsize);
1455 if (_PyBytes_Resize(&ret, bufsize) < 0) {
1456 BZ2_bzCompressEnd(bzs);
1457 goto error;
1458 }
1459 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1460 - totalout);
1461 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1462 }
1463 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001464
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001465 if (_PyBytes_Resize(&ret,
1466 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1467 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001468
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001469 RELEASE_LOCK(self);
1470 PyBuffer_Release(&pdata);
1471 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001472
1473error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001474 RELEASE_LOCK(self);
1475 PyBuffer_Release(&pdata);
1476 Py_XDECREF(ret);
1477 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001478}
1479
1480PyDoc_STRVAR(BZ2Comp_flush__doc__,
1481"flush() -> string\n\
1482\n\
1483Finish the compression process and return what is left in internal buffers.\n\
1484You must not use the compressor object after calling this method.\n\
1485");
1486
1487static PyObject *
1488BZ2Comp_flush(BZ2CompObject *self)
1489{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001490 int bufsize = SMALLCHUNK;
1491 PyObject *ret = NULL;
1492 bz_stream *bzs = &self->bzs;
1493 PY_LONG_LONG totalout;
1494 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001495
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001496 ACQUIRE_LOCK(self);
1497 if (!self->running) {
1498 PyErr_SetString(PyExc_ValueError, "object was already "
1499 "flushed");
1500 goto error;
1501 }
1502 self->running = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001503
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001504 ret = PyBytes_FromStringAndSize(NULL, bufsize);
1505 if (!ret)
1506 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001507
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001508 bzs->next_out = BUF(ret);
1509 bzs->avail_out = bufsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001510
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001511 totalout = BZS_TOTAL_OUT(bzs);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001512
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001513 for (;;) {
1514 Py_BEGIN_ALLOW_THREADS
1515 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1516 Py_END_ALLOW_THREADS
1517 if (bzerror == BZ_STREAM_END) {
1518 break;
1519 } else if (bzerror != BZ_FINISH_OK) {
1520 Util_CatchBZ2Error(bzerror);
1521 goto error;
1522 }
1523 if (bzs->avail_out == 0) {
1524 bufsize = Util_NewBufferSize(bufsize);
1525 if (_PyBytes_Resize(&ret, bufsize) < 0)
1526 goto error;
1527 bzs->next_out = BUF(ret);
1528 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1529 - totalout);
1530 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1531 }
1532 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001533
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001534 if (bzs->avail_out != 0) {
1535 if (_PyBytes_Resize(&ret,
1536 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1537 goto error;
1538 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001539
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001540 RELEASE_LOCK(self);
1541 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001542
1543error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001544 RELEASE_LOCK(self);
1545 Py_XDECREF(ret);
1546 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001547}
1548
1549static PyMethodDef BZ2Comp_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001550 {"compress", (PyCFunction)BZ2Comp_compress, METH_VARARGS,
1551 BZ2Comp_compress__doc__},
1552 {"flush", (PyCFunction)BZ2Comp_flush, METH_NOARGS,
1553 BZ2Comp_flush__doc__},
1554 {NULL, NULL} /* sentinel */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001555};
1556
1557
1558/* ===================================================================== */
1559/* Slot definitions for BZ2Comp_Type. */
1560
1561static int
1562BZ2Comp_init(BZ2CompObject *self, PyObject *args, PyObject *kwargs)
1563{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001564 int compresslevel = 9;
1565 int bzerror;
1566 static char *kwlist[] = {"compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001567
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001568 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:BZ2Compressor",
1569 kwlist, &compresslevel))
1570 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001571
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001572 if (compresslevel < 1 || compresslevel > 9) {
1573 PyErr_SetString(PyExc_ValueError,
1574 "compresslevel must be between 1 and 9");
1575 goto error;
1576 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001577
1578#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001579 self->lock = PyThread_allocate_lock();
1580 if (!self->lock) {
1581 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1582 goto error;
1583 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001584#endif
1585
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001586 memset(&self->bzs, 0, sizeof(bz_stream));
1587 bzerror = BZ2_bzCompressInit(&self->bzs, compresslevel, 0, 0);
1588 if (bzerror != BZ_OK) {
1589 Util_CatchBZ2Error(bzerror);
1590 goto error;
1591 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001592
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001593 self->running = 1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001594
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001595 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001596error:
1597#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001598 if (self->lock) {
1599 PyThread_free_lock(self->lock);
1600 self->lock = NULL;
1601 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001602#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001603 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001604}
1605
1606static void
1607BZ2Comp_dealloc(BZ2CompObject *self)
1608{
1609#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001610 if (self->lock)
1611 PyThread_free_lock(self->lock);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001612#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001613 BZ2_bzCompressEnd(&self->bzs);
1614 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001615}
1616
1617
1618/* ===================================================================== */
1619/* BZ2Comp_Type definition. */
1620
1621PyDoc_STRVAR(BZ2Comp__doc__,
1622"BZ2Compressor([compresslevel=9]) -> compressor object\n\
1623\n\
1624Create a new compressor object. This object may be used to compress\n\
1625data sequentially. If you want to compress data in one shot, use the\n\
1626compress() function instead. The compresslevel parameter, if given,\n\
1627must be a number between 1 and 9.\n\
1628");
1629
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001630static PyTypeObject BZ2Comp_Type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001631 PyVarObject_HEAD_INIT(NULL, 0)
1632 "bz2.BZ2Compressor", /*tp_name*/
1633 sizeof(BZ2CompObject), /*tp_basicsize*/
1634 0, /*tp_itemsize*/
1635 (destructor)BZ2Comp_dealloc, /*tp_dealloc*/
1636 0, /*tp_print*/
1637 0, /*tp_getattr*/
1638 0, /*tp_setattr*/
1639 0, /*tp_reserved*/
1640 0, /*tp_repr*/
1641 0, /*tp_as_number*/
1642 0, /*tp_as_sequence*/
1643 0, /*tp_as_mapping*/
1644 0, /*tp_hash*/
1645 0, /*tp_call*/
1646 0, /*tp_str*/
1647 PyObject_GenericGetAttr,/*tp_getattro*/
1648 PyObject_GenericSetAttr,/*tp_setattro*/
1649 0, /*tp_as_buffer*/
1650 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1651 BZ2Comp__doc__, /*tp_doc*/
1652 0, /*tp_traverse*/
1653 0, /*tp_clear*/
1654 0, /*tp_richcompare*/
1655 0, /*tp_weaklistoffset*/
1656 0, /*tp_iter*/
1657 0, /*tp_iternext*/
1658 BZ2Comp_methods, /*tp_methods*/
1659 0, /*tp_members*/
1660 0, /*tp_getset*/
1661 0, /*tp_base*/
1662 0, /*tp_dict*/
1663 0, /*tp_descr_get*/
1664 0, /*tp_descr_set*/
1665 0, /*tp_dictoffset*/
1666 (initproc)BZ2Comp_init, /*tp_init*/
1667 PyType_GenericAlloc, /*tp_alloc*/
1668 PyType_GenericNew, /*tp_new*/
1669 PyObject_Free, /*tp_free*/
1670 0, /*tp_is_gc*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001671};
1672
1673
1674/* ===================================================================== */
1675/* Members of BZ2Decomp. */
1676
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001677#undef OFF
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001678#define OFF(x) offsetof(BZ2DecompObject, x)
1679
1680static PyMemberDef BZ2Decomp_members[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001681 {"unused_data", T_OBJECT, OFF(unused_data), READONLY},
1682 {NULL} /* Sentinel */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001683};
1684
1685
1686/* ===================================================================== */
1687/* Methods of BZ2Decomp. */
1688
1689PyDoc_STRVAR(BZ2Decomp_decompress__doc__,
1690"decompress(data) -> string\n\
1691\n\
1692Provide more data to the decompressor object. It will return chunks\n\
1693of decompressed data whenever possible. If you try to decompress data\n\
1694after the end of stream is found, EOFError will be raised. If any data\n\
1695was found after the end of stream, it'll be ignored and saved in\n\
1696unused_data attribute.\n\
1697");
1698
1699static PyObject *
1700BZ2Decomp_decompress(BZ2DecompObject *self, PyObject *args)
1701{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001702 Py_buffer pdata;
1703 char *data;
1704 int datasize;
1705 int bufsize = SMALLCHUNK;
1706 PY_LONG_LONG totalout;
1707 PyObject *ret = NULL;
1708 bz_stream *bzs = &self->bzs;
1709 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001710
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001711 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
1712 return NULL;
1713 data = pdata.buf;
1714 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001715
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001716 ACQUIRE_LOCK(self);
1717 if (!self->running) {
1718 PyErr_SetString(PyExc_EOFError, "end of stream was "
1719 "already found");
1720 goto error;
1721 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001722
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001723 ret = PyBytes_FromStringAndSize(NULL, bufsize);
1724 if (!ret)
1725 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001726
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001727 bzs->next_in = data;
1728 bzs->avail_in = datasize;
1729 bzs->next_out = BUF(ret);
1730 bzs->avail_out = bufsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001731
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001732 totalout = BZS_TOTAL_OUT(bzs);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001733
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001734 for (;;) {
1735 Py_BEGIN_ALLOW_THREADS
1736 bzerror = BZ2_bzDecompress(bzs);
1737 Py_END_ALLOW_THREADS
1738 if (bzerror == BZ_STREAM_END) {
1739 if (bzs->avail_in != 0) {
1740 Py_DECREF(self->unused_data);
1741 self->unused_data =
1742 PyBytes_FromStringAndSize(bzs->next_in,
1743 bzs->avail_in);
1744 }
1745 self->running = 0;
1746 break;
1747 }
1748 if (bzerror != BZ_OK) {
1749 Util_CatchBZ2Error(bzerror);
1750 goto error;
1751 }
1752 if (bzs->avail_in == 0)
1753 break; /* no more input data */
1754 if (bzs->avail_out == 0) {
1755 bufsize = Util_NewBufferSize(bufsize);
1756 if (_PyBytes_Resize(&ret, bufsize) < 0) {
1757 BZ2_bzDecompressEnd(bzs);
1758 goto error;
1759 }
1760 bzs->next_out = BUF(ret);
1761 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1762 - totalout);
1763 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1764 }
1765 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001766
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001767 if (bzs->avail_out != 0) {
1768 if (_PyBytes_Resize(&ret,
1769 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1770 goto error;
1771 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001772
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001773 RELEASE_LOCK(self);
1774 PyBuffer_Release(&pdata);
1775 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001776
1777error:
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001778 RELEASE_LOCK(self);
1779 PyBuffer_Release(&pdata);
1780 Py_XDECREF(ret);
1781 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001782}
1783
1784static PyMethodDef BZ2Decomp_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001785 {"decompress", (PyCFunction)BZ2Decomp_decompress, METH_VARARGS, BZ2Decomp_decompress__doc__},
1786 {NULL, NULL} /* sentinel */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001787};
1788
1789
1790/* ===================================================================== */
1791/* Slot definitions for BZ2Decomp_Type. */
1792
1793static int
1794BZ2Decomp_init(BZ2DecompObject *self, PyObject *args, PyObject *kwargs)
1795{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001796 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001797
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001798 if (!PyArg_ParseTuple(args, ":BZ2Decompressor"))
1799 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001800
1801#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001802 self->lock = PyThread_allocate_lock();
1803 if (!self->lock) {
1804 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1805 goto error;
1806 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001807#endif
1808
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001809 self->unused_data = PyBytes_FromStringAndSize("", 0);
1810 if (!self->unused_data)
1811 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001812
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001813 memset(&self->bzs, 0, sizeof(bz_stream));
1814 bzerror = BZ2_bzDecompressInit(&self->bzs, 0, 0);
1815 if (bzerror != BZ_OK) {
1816 Util_CatchBZ2Error(bzerror);
1817 goto error;
1818 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001819
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001820 self->running = 1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001821
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001822 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001823
1824error:
1825#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001826 if (self->lock) {
1827 PyThread_free_lock(self->lock);
1828 self->lock = NULL;
1829 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001830#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001831 Py_CLEAR(self->unused_data);
1832 return -1;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001833}
1834
1835static void
1836BZ2Decomp_dealloc(BZ2DecompObject *self)
1837{
1838#ifdef WITH_THREAD
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001839 if (self->lock)
1840 PyThread_free_lock(self->lock);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001841#endif
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001842 Py_XDECREF(self->unused_data);
1843 BZ2_bzDecompressEnd(&self->bzs);
1844 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001845}
1846
1847
1848/* ===================================================================== */
1849/* BZ2Decomp_Type definition. */
1850
1851PyDoc_STRVAR(BZ2Decomp__doc__,
1852"BZ2Decompressor() -> decompressor object\n\
1853\n\
1854Create a new decompressor object. This object may be used to decompress\n\
1855data sequentially. If you want to decompress data in one shot, use the\n\
1856decompress() function instead.\n\
1857");
1858
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001859static PyTypeObject BZ2Decomp_Type = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001860 PyVarObject_HEAD_INIT(NULL, 0)
1861 "bz2.BZ2Decompressor", /*tp_name*/
1862 sizeof(BZ2DecompObject), /*tp_basicsize*/
1863 0, /*tp_itemsize*/
1864 (destructor)BZ2Decomp_dealloc, /*tp_dealloc*/
1865 0, /*tp_print*/
1866 0, /*tp_getattr*/
1867 0, /*tp_setattr*/
1868 0, /*tp_reserved*/
1869 0, /*tp_repr*/
1870 0, /*tp_as_number*/
1871 0, /*tp_as_sequence*/
1872 0, /*tp_as_mapping*/
1873 0, /*tp_hash*/
1874 0, /*tp_call*/
1875 0, /*tp_str*/
1876 PyObject_GenericGetAttr,/*tp_getattro*/
1877 PyObject_GenericSetAttr,/*tp_setattro*/
1878 0, /*tp_as_buffer*/
1879 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1880 BZ2Decomp__doc__, /*tp_doc*/
1881 0, /*tp_traverse*/
1882 0, /*tp_clear*/
1883 0, /*tp_richcompare*/
1884 0, /*tp_weaklistoffset*/
1885 0, /*tp_iter*/
1886 0, /*tp_iternext*/
1887 BZ2Decomp_methods, /*tp_methods*/
1888 BZ2Decomp_members, /*tp_members*/
1889 0, /*tp_getset*/
1890 0, /*tp_base*/
1891 0, /*tp_dict*/
1892 0, /*tp_descr_get*/
1893 0, /*tp_descr_set*/
1894 0, /*tp_dictoffset*/
1895 (initproc)BZ2Decomp_init, /*tp_init*/
1896 PyType_GenericAlloc, /*tp_alloc*/
1897 PyType_GenericNew, /*tp_new*/
1898 PyObject_Free, /*tp_free*/
1899 0, /*tp_is_gc*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001900};
1901
1902
1903/* ===================================================================== */
1904/* Module functions. */
1905
1906PyDoc_STRVAR(bz2_compress__doc__,
1907"compress(data [, compresslevel=9]) -> string\n\
1908\n\
1909Compress data in one shot. If you want to compress data sequentially,\n\
1910use an instance of BZ2Compressor instead. The compresslevel parameter, if\n\
1911given, must be a number between 1 and 9.\n\
1912");
1913
1914static PyObject *
1915bz2_compress(PyObject *self, PyObject *args, PyObject *kwargs)
1916{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001917 int compresslevel=9;
1918 Py_buffer pdata;
1919 char *data;
1920 int datasize;
1921 int bufsize;
1922 PyObject *ret = NULL;
1923 bz_stream _bzs;
1924 bz_stream *bzs = &_bzs;
1925 int bzerror;
1926 static char *kwlist[] = {"data", "compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001927
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001928 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|i",
1929 kwlist, &pdata,
1930 &compresslevel))
1931 return NULL;
1932 data = pdata.buf;
1933 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001934
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001935 if (compresslevel < 1 || compresslevel > 9) {
1936 PyErr_SetString(PyExc_ValueError,
1937 "compresslevel must be between 1 and 9");
1938 PyBuffer_Release(&pdata);
1939 return NULL;
1940 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001941
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001942 /* Conforming to bz2 manual, this is large enough to fit compressed
1943 * data in one shot. We will check it later anyway. */
1944 bufsize = datasize + (datasize/100+1) + 600;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001945
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001946 ret = PyBytes_FromStringAndSize(NULL, bufsize);
1947 if (!ret) {
1948 PyBuffer_Release(&pdata);
1949 return NULL;
1950 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001951
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001952 memset(bzs, 0, sizeof(bz_stream));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001953
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001954 bzs->next_in = data;
1955 bzs->avail_in = datasize;
1956 bzs->next_out = BUF(ret);
1957 bzs->avail_out = bufsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001958
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001959 bzerror = BZ2_bzCompressInit(bzs, compresslevel, 0, 0);
1960 if (bzerror != BZ_OK) {
1961 Util_CatchBZ2Error(bzerror);
1962 PyBuffer_Release(&pdata);
1963 Py_DECREF(ret);
1964 return NULL;
1965 }
Tim Peterse3228092002-11-09 04:21:44 +00001966
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001967 for (;;) {
1968 Py_BEGIN_ALLOW_THREADS
1969 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1970 Py_END_ALLOW_THREADS
1971 if (bzerror == BZ_STREAM_END) {
1972 break;
1973 } else if (bzerror != BZ_FINISH_OK) {
1974 BZ2_bzCompressEnd(bzs);
1975 Util_CatchBZ2Error(bzerror);
1976 PyBuffer_Release(&pdata);
1977 Py_DECREF(ret);
1978 return NULL;
1979 }
1980 if (bzs->avail_out == 0) {
1981 bufsize = Util_NewBufferSize(bufsize);
1982 if (_PyBytes_Resize(&ret, bufsize) < 0) {
1983 BZ2_bzCompressEnd(bzs);
1984 PyBuffer_Release(&pdata);
1985 return NULL;
1986 }
1987 bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
1988 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1989 }
1990 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001991
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001992 if (bzs->avail_out != 0) {
1993 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
1994 ret = NULL;
1995 }
1996 }
1997 BZ2_bzCompressEnd(bzs);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001998
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00001999 PyBuffer_Release(&pdata);
2000 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002001}
2002
2003PyDoc_STRVAR(bz2_decompress__doc__,
2004"decompress(data) -> decompressed data\n\
2005\n\
2006Decompress data in one shot. If you want to decompress data sequentially,\n\
2007use an instance of BZ2Decompressor instead.\n\
2008");
2009
2010static PyObject *
2011bz2_decompress(PyObject *self, PyObject *args)
2012{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002013 Py_buffer pdata;
2014 char *data;
2015 int datasize;
2016 int bufsize = SMALLCHUNK;
2017 PyObject *ret;
2018 bz_stream _bzs;
2019 bz_stream *bzs = &_bzs;
2020 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002021
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002022 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
2023 return NULL;
2024 data = pdata.buf;
2025 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002026
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002027 if (datasize == 0) {
2028 PyBuffer_Release(&pdata);
2029 return PyBytes_FromStringAndSize("", 0);
2030 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002031
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002032 ret = PyBytes_FromStringAndSize(NULL, bufsize);
2033 if (!ret) {
2034 PyBuffer_Release(&pdata);
2035 return NULL;
2036 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002037
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002038 memset(bzs, 0, sizeof(bz_stream));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002039
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002040 bzs->next_in = data;
2041 bzs->avail_in = datasize;
2042 bzs->next_out = BUF(ret);
2043 bzs->avail_out = bufsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002044
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002045 bzerror = BZ2_bzDecompressInit(bzs, 0, 0);
2046 if (bzerror != BZ_OK) {
2047 Util_CatchBZ2Error(bzerror);
2048 Py_DECREF(ret);
2049 PyBuffer_Release(&pdata);
2050 return NULL;
2051 }
Tim Peterse3228092002-11-09 04:21:44 +00002052
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002053 for (;;) {
2054 Py_BEGIN_ALLOW_THREADS
2055 bzerror = BZ2_bzDecompress(bzs);
2056 Py_END_ALLOW_THREADS
2057 if (bzerror == BZ_STREAM_END) {
2058 break;
2059 } else if (bzerror != BZ_OK) {
2060 BZ2_bzDecompressEnd(bzs);
2061 Util_CatchBZ2Error(bzerror);
2062 PyBuffer_Release(&pdata);
2063 Py_DECREF(ret);
2064 return NULL;
2065 }
2066 if (bzs->avail_in == 0) {
2067 BZ2_bzDecompressEnd(bzs);
2068 PyErr_SetString(PyExc_ValueError,
2069 "couldn't find end of stream");
2070 PyBuffer_Release(&pdata);
2071 Py_DECREF(ret);
2072 return NULL;
2073 }
2074 if (bzs->avail_out == 0) {
2075 bufsize = Util_NewBufferSize(bufsize);
2076 if (_PyBytes_Resize(&ret, bufsize) < 0) {
2077 BZ2_bzDecompressEnd(bzs);
2078 PyBuffer_Release(&pdata);
2079 return NULL;
2080 }
2081 bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
2082 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
2083 }
2084 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002085
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002086 if (bzs->avail_out != 0) {
2087 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
2088 ret = NULL;
2089 }
2090 }
2091 BZ2_bzDecompressEnd(bzs);
2092 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002093
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002094 return ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002095}
2096
2097static PyMethodDef bz2_methods[] = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002098 {"compress", (PyCFunction) bz2_compress, METH_VARARGS|METH_KEYWORDS,
2099 bz2_compress__doc__},
2100 {"decompress", (PyCFunction) bz2_decompress, METH_VARARGS,
2101 bz2_decompress__doc__},
2102 {NULL, NULL} /* sentinel */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002103};
2104
2105/* ===================================================================== */
2106/* Initialization function. */
2107
2108PyDoc_STRVAR(bz2__doc__,
2109"The python bz2 module provides a comprehensive interface for\n\
2110the bz2 compression library. It implements a complete file\n\
2111interface, one shot (de)compression functions, and types for\n\
2112sequential (de)compression.\n\
2113");
2114
Martin v. Löwis1a214512008-06-11 05:26:20 +00002115
2116static struct PyModuleDef bz2module = {
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002117 PyModuleDef_HEAD_INIT,
2118 "bz2",
2119 bz2__doc__,
2120 -1,
2121 bz2_methods,
2122 NULL,
2123 NULL,
2124 NULL,
2125 NULL
Martin v. Löwis1a214512008-06-11 05:26:20 +00002126};
2127
Neal Norwitz21d896c2003-07-01 20:15:21 +00002128PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00002129PyInit_bz2(void)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002130{
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002131 PyObject *m;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002132
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002133 Py_TYPE(&BZ2File_Type) = &PyType_Type;
2134 Py_TYPE(&BZ2Comp_Type) = &PyType_Type;
2135 Py_TYPE(&BZ2Decomp_Type) = &PyType_Type;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002136
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002137 m = PyModule_Create(&bz2module);
2138 if (m == NULL)
2139 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002140
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002141 PyModule_AddObject(m, "__author__", PyUnicode_FromString(__author__));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002142
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002143 Py_INCREF(&BZ2File_Type);
2144 PyModule_AddObject(m, "BZ2File", (PyObject *)&BZ2File_Type);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002145
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002146 Py_INCREF(&BZ2Comp_Type);
2147 PyModule_AddObject(m, "BZ2Compressor", (PyObject *)&BZ2Comp_Type);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002148
Antoine Pitrou7f14f0d2010-05-09 16:14:21 +00002149 Py_INCREF(&BZ2Decomp_Type);
2150 PyModule_AddObject(m, "BZ2Decompressor", (PyObject *)&BZ2Decomp_Type);
2151 return m;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002152}