blob: 3e8eb8ba4010d33e213a756ac133a776207545b3 [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
Christian Heimes90aa7642007-12-19 02:45:37 +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) \
51 (((long)bzs->total_out_hi32 << 32) + bzs->total_out_lo32)
52#elif SIZEOF_LONG_LONG >= 8
53#define BZS_TOTAL_OUT(bzs) \
Martin v. Löwisb9a0f912003-03-29 10:06:18 +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) \
Neal Norwitz20bad742006-01-17 05:27:39 +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 { \
82 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 */
94#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 {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000103 PyObject_HEAD
Guido van Rossumf09ca142007-06-13 00:03:05 +0000104 FILE *rawfp;
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000105
106 char* f_buf; /* Allocated readahead buffer */
107 char* f_bufend; /* Points after last occupied position */
108 char* f_bufptr; /* Current buffer position */
109
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000110 BZFILE *fp;
111 int mode;
Georg Brandla8bcecc2005-09-03 07:49:53 +0000112 Py_off_t pos;
113 Py_off_t size;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000114#ifdef WITH_THREAD
115 PyThread_type_lock lock;
116#endif
117} BZ2FileObject;
118
119typedef struct {
120 PyObject_HEAD
121 bz_stream bzs;
122 int running;
123#ifdef WITH_THREAD
124 PyThread_type_lock lock;
125#endif
126} BZ2CompObject;
127
128typedef struct {
129 PyObject_HEAD
130 bz_stream bzs;
131 int running;
132 PyObject *unused_data;
133#ifdef WITH_THREAD
134 PyThread_type_lock lock;
135#endif
136} BZ2DecompObject;
137
138/* ===================================================================== */
139/* Utility functions. */
140
141static int
142Util_CatchBZ2Error(int bzerror)
143{
144 int ret = 0;
145 switch(bzerror) {
146 case BZ_OK:
147 case BZ_STREAM_END:
148 break;
149
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +0000150#ifdef BZ_CONFIG_ERROR
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000166 case BZ_MEM_ERROR:
167 PyErr_NoMemory();
168 ret = 1;
169 break;
170
171 case BZ_DATA_ERROR:
172 case BZ_DATA_ERROR_MAGIC:
173 PyErr_SetString(PyExc_IOError, "invalid data stream");
174 ret = 1;
175 break;
176
177 case BZ_IO_ERROR:
178 PyErr_SetString(PyExc_IOError, "unknown IO error");
179 ret = 1;
180 break;
181
182 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;
188
189 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;
197}
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{
215 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;
224}
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{
230 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;
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000237 int bytes_read;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000238
239 total_v_size = n > 0 ? n : 100;
Christian Heimes72b710a2008-05-26 13:28:38 +0000240 v = PyBytes_FromStringAndSize((char *)NULL, total_v_size);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000241 if (v == NULL)
242 return NULL;
243
244 buf = BUF(v);
245 end = buf + total_v_size;
246
247 for (;;) {
248 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000249 do {
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000250 bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
Guido van Rossumf09ca142007-06-13 00:03:05 +0000251 f->pos++;
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000252 if (bytes_read == 0)
253 break;
Guido van Rossumf09ca142007-06-13 00:03:05 +0000254 *buf++ = c;
255 } while (bzerror == BZ_OK && c != '\n' && buf != end);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000256 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000257 if (bzerror == BZ_STREAM_END) {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000258 f->size = f->pos;
259 f->mode = MODE_READ_EOF;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000260 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 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000280 if (_PyBytes_Resize(&v, total_v_size) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000281 return NULL;
Guido van Rossum522a6c62007-05-22 23:13:45 +0000282 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000283 buf = BUF(v) + used_v_size;
284 end = BUF(v) + total_v_size;
285 }
286
287 used_v_size = buf - BUF(v);
Guido van Rossum522a6c62007-05-22 23:13:45 +0000288 if (used_v_size != total_v_size) {
Christian Heimes72b710a2008-05-26 13:28:38 +0000289 if (_PyBytes_Resize(&v, used_v_size) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000290 v = NULL;
291 }
292 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000293 return v;
294}
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{
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000300 if (f->f_buf != NULL) {
301 PyMem_Free(f->f_buf);
302 f->f_buf = NULL;
303 }
304}
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{
310 int chunksize;
311 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000312
313 if (f->f_buf != NULL) {
Tim Peterse3228092002-11-09 04:21:44 +0000314 if((f->f_bufend - f->f_bufptr) >= 1)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000315 return 0;
316 else
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000317 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000318 }
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000319 if (f->mode == MODE_READ_EOF) {
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000320 f->f_bufptr = f->f_buf;
321 f->f_bufend = f->f_buf;
322 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000323 }
324 if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {
Georg Brandl6aa2d1f2008-08-12 08:35:52 +0000325 PyErr_NoMemory();
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000326 return -1;
327 }
328 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000329 chunksize = BZ2_bzRead(&bzerror, f->fp, f->f_buf, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000330 Py_END_ALLOW_THREADS
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000331 f->pos += chunksize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000332 if (bzerror == BZ_STREAM_END) {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000333 f->size = f->pos;
334 f->mode = MODE_READ_EOF;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000335 } else if (bzerror != BZ_OK) {
336 Util_CatchBZ2Error(bzerror);
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000337 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000338 return -1;
339 }
340 f->f_bufptr = f->f_buf;
341 f->f_bufend = f->f_buf + chunksize;
342 return 0;
343}
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{
Christian Heimes72b710a2008-05-26 13:28:38 +0000350 PyBytesObject* s;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000351 char *bufptr;
352 char *buf;
353 int len;
354
355 if (f->f_buf == NULL)
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000356 if (Util_ReadAhead(f, bufsize) < 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000357 return NULL;
358
359 len = f->f_bufend - f->f_bufptr;
Tim Peterse3228092002-11-09 04:21:44 +0000360 if (len == 0)
Christian Heimes72b710a2008-05-26 13:28:38 +0000361 return (PyBytesObject *)
362 PyBytes_FromStringAndSize(NULL, skip);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000363 bufptr = memchr(f->f_bufptr, '\n', len);
364 if (bufptr != NULL) {
365 bufptr++; /* Count the '\n' */
366 len = bufptr - f->f_bufptr;
Christian Heimes72b710a2008-05-26 13:28:38 +0000367 s = (PyBytesObject *)
368 PyBytes_FromStringAndSize(NULL, skip+len);
Tim Peterse3228092002-11-09 04:21:44 +0000369 if (s == NULL)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000370 return NULL;
Christian Heimes72b710a2008-05-26 13:28:38 +0000371 memcpy(PyBytes_AS_STRING(s)+skip, f->f_bufptr, len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000372 f->f_bufptr = bufptr;
373 if (bufptr == f->f_bufend)
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000374 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000375 } else {
376 bufptr = f->f_bufptr;
377 buf = f->f_buf;
378 f->f_buf = NULL; /* Force new readahead buffer */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000379 s = Util_ReadAheadGetLineSkip(f, skip+len,
380 bufsize + (bufsize>>2));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000381 if (s == NULL) {
382 PyMem_Free(buf);
383 return NULL;
384 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000385 memcpy(PyBytes_AS_STRING(s)+skip, bufptr, len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000386 PyMem_Free(buf);
387 }
388 return s;
389}
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{
405 long bytesrequested = -1;
406 size_t bytesread, buffersize, chunksize;
407 int bzerror;
408 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000409
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000410 if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))
411 return NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000412
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000413 ACQUIRE_LOCK(self);
414 switch (self->mode) {
415 case MODE_READ:
416 break;
417 case MODE_READ_EOF:
Christian Heimes72b710a2008-05-26 13:28:38 +0000418 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000419 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 }
429
430 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,
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +0000436 "requested number of bytes is "
437 "more than a Python string can hold");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000438 goto cleanup;
439 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000440 ret = PyBytes_FromStringAndSize((char *)NULL, buffersize);
Guido van Rossum75c26bc2007-08-07 23:29:20 +0000441 if (ret == NULL || buffersize == 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000442 goto cleanup;
443 bytesread = 0;
444
445 for (;;) {
446 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000447 chunksize = BZ2_bzRead(&bzerror, self->fp,
448 BUF(ret)+bytesread,
449 buffersize-bytesread);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000450 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);
Christian Heimes72b710a2008-05-26 13:28:38 +0000465 if (_PyBytes_Resize(&ret, buffersize) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000466 ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000467 goto cleanup;
Guido van Rossum522a6c62007-05-22 23:13:45 +0000468 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000469 } else {
470 break;
471 }
472 }
Guido van Rossum522a6c62007-05-22 23:13:45 +0000473 if (bytesread != buffersize) {
Christian Heimes72b710a2008-05-26 13:28:38 +0000474 if (_PyBytes_Resize(&ret, bytesread) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000475 ret = NULL;
476 }
477 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000478
479cleanup:
480 RELEASE_LOCK(self);
481 return ret;
482}
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{
496 PyObject *ret = NULL;
497 int sizehint = -1;
498
499 if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))
500 return NULL;
501
502 ACQUIRE_LOCK(self);
503 switch (self->mode) {
504 case MODE_READ:
505 break;
506 case MODE_READ_EOF:
Christian Heimes72b710a2008-05-26 13:28:38 +0000507 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000508 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 }
518
519 if (sizehint == 0)
Christian Heimes72b710a2008-05-26 13:28:38 +0000520 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000521 else
522 ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);
523
524cleanup:
525 RELEASE_LOCK(self);
526 return ret;
527}
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{
541 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;
555
556 if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))
557 return NULL;
558
559 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 }
575
576 if ((list = PyList_New(0)) == NULL)
577 goto cleanup;
578
579 for (;;) {
580 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000581 nread = BZ2_bzRead(&bzerror, self->fp,
582 buffer+nfilled, buffersize-nfilled);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000583 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);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000602 if (!shortread && p == NULL) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000603 /* Need a larger buffer to fit this line */
604 nfilled += nread;
605 buffersize *= 2;
606 if (buffersize > INT_MAX) {
607 PyErr_SetString(PyExc_OverflowError,
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000608 "line is longer than a Python string can hold");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000609 goto error;
610 }
611 if (big_buffer == NULL) {
612 /* Create the big buffer */
Christian Heimes72b710a2008-05-26 13:28:38 +0000613 big_buffer = PyBytes_FromStringAndSize(
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000614 NULL, buffersize);
615 if (big_buffer == NULL)
616 goto error;
Christian Heimes72b710a2008-05-26 13:28:38 +0000617 buffer = PyBytes_AS_STRING(big_buffer);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000618 memcpy(buffer, small_buffer, nfilled);
619 }
620 else {
621 /* Grow the big buffer */
Christian Heimes72b710a2008-05-26 13:28:38 +0000622 if (_PyBytes_Resize(&big_buffer, buffersize) < 0){
Guido van Rossum522a6c62007-05-22 23:13:45 +0000623 big_buffer = NULL;
624 goto error;
625 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000626 buffer = PyBytes_AS_STRING(big_buffer);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000627 }
Guido van Rossum522a6c62007-05-22 23:13:45 +0000628 continue;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000629 }
630 end = buffer+nfilled+nread;
631 q = buffer;
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000632 while (p != NULL) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000633 /* Process complete lines */
634 p++;
Christian Heimes72b710a2008-05-26 13:28:38 +0000635 line = PyBytes_FromStringAndSize(q, p-q);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000636 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);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000644 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000645 /* 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 */
Christian Heimes72b710a2008-05-26 13:28:38 +0000658 line = PyBytes_FromStringAndSize(buffer, nfilled);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000659 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 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000668 PyBytes_Concat(&line, rest);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000669 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 }
678
679 cleanup:
680 RELEASE_LOCK(self);
681 if (big_buffer) {
682 Py_DECREF(big_buffer);
683 }
684 return list;
685}
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{
698 PyObject *ret = NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +0000699 Py_buffer pbuf;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000700 char *buf;
701 int len;
702 int bzerror;
703
Martin v. Löwis423be952008-08-13 15:53:07 +0000704 if (!PyArg_ParseTuple(args, "y*:write", &pbuf))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000705 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +0000706 buf = pbuf.buf;
707 len = pbuf.len;
Tim Peterse3228092002-11-09 04:21:44 +0000708
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000709 ACQUIRE_LOCK(self);
710 switch (self->mode) {
711 case MODE_WRITE:
712 break;
Tim Peterse3228092002-11-09 04:21:44 +0000713
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000714 case MODE_CLOSED:
715 PyErr_SetString(PyExc_ValueError,
716 "I/O operation on closed file");
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000717 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000718
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000719 default:
720 PyErr_SetString(PyExc_IOError,
721 "file is not ready for writing");
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000722 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000723 }
724
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000730 if (bzerror != BZ_OK) {
731 Util_CatchBZ2Error(bzerror);
732 goto cleanup;
733 }
Tim Peterse3228092002-11-09 04:21:44 +0000734
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000735 Py_INCREF(Py_None);
736 ret = Py_None;
737
738cleanup:
Martin v. Löwis423be952008-08-13 15:53:07 +0000739 PyBuffer_Release(&pbuf);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000740 RELEASE_LOCK(self);
741 return ret;
742}
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
757 PyObject *list = NULL;
758 PyObject *iter = NULL;
759 PyObject *ret = NULL;
760 PyObject *line;
761 int i, j, index, len, islist;
762 int bzerror;
763
764 ACQUIRE_LOCK(self);
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000765 switch (self->mode) {
766 case MODE_WRITE:
767 break;
768
769 case MODE_CLOSED:
770 PyErr_SetString(PyExc_ValueError,
771 "I/O operation on closed file");
772 goto error;
773
774 default:
775 PyErr_SetString(PyExc_IOError,
776 "file is not ready for writing");
777 goto error;
778 }
779
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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 }
792
793 /* 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;
817
Guido van Rossum522a6c62007-05-22 23:13:45 +0000818 /* Check that all entries are indeed byte strings. If not,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000819 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);
Christian Heimes72b710a2008-05-26 13:28:38 +0000825 if (!PyBytes_Check(v)) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000826 const char *buffer;
Martin v. Löwis18e16552006-02-15 17:27:45 +0000827 Py_ssize_t len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000828 if (PyObject_AsCharBuffer(v, &buffer, &len)) {
829 PyErr_SetString(PyExc_TypeError,
830 "writelines() "
831 "argument must be "
832 "a sequence of "
Guido van Rossum522a6c62007-05-22 23:13:45 +0000833 "bytes objects");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000834 goto error;
835 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000836 line = PyBytes_FromStringAndSize(buffer,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000837 len);
838 if (line == NULL)
839 goto error;
840 Py_DECREF(v);
841 PyList_SET_ITEM(list, i, line);
842 }
843 }
844
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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);
Christian Heimes72b710a2008-05-26 13:28:38 +0000850 len = PyBytes_GET_SIZE(line);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000851 BZ2_bzWrite (&bzerror, self->fp,
Christian Heimes72b710a2008-05-26 13:28:38 +0000852 PyBytes_AS_STRING(line), len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000853 if (bzerror != BZ_OK) {
854 Py_BLOCK_THREADS
855 Util_CatchBZ2Error(bzerror);
856 goto error;
857 }
858 }
859 Py_END_ALLOW_THREADS
860
861 if (j < CHUNKSIZE)
862 break;
863 }
864
865 Py_INCREF(Py_None);
866 ret = Py_None;
867
868 error:
869 RELEASE_LOCK(self);
870 Py_XDECREF(list);
871 Py_XDECREF(iter);
872 return ret;
873#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{
892 int where = 0;
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000893 PyObject *offobj;
894 Py_off_t offset;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000895 char small_buffer[SMALLCHUNK];
896 char *buffer = small_buffer;
897 size_t buffersize = SMALLCHUNK;
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000898 Py_off_t bytesread = 0;
Georg Brandla8bcecc2005-09-03 07:49:53 +0000899 size_t readsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000900 int chunksize;
901 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000902 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000903
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000904 if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))
905 return NULL;
906#if !defined(HAVE_LARGEFILE_SUPPORT)
Christian Heimes217cfd12007-12-02 14:31:20 +0000907 offset = PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000908#else
909 offset = PyLong_Check(offobj) ?
Christian Heimes217cfd12007-12-02 14:31:20 +0000910 PyLong_AsLongLong(offobj) : PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000911#endif
912 if (PyErr_Occurred())
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000913 return NULL;
914
915 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
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000922 case MODE_CLOSED:
923 PyErr_SetString(PyExc_ValueError,
924 "I/O operation on closed file");
Thomas Wouters89f507f2006-12-13 04:49:30 +0000925 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000926
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000927 default:
928 PyErr_SetString(PyExc_IOError,
929 "seek works only while reading");
Thomas Wouters89f507f2006-12-13 04:49:30 +0000930 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000931 }
932
Georg Brandl47fab922006-02-18 21:57:25 +0000933 if (where == 2) {
934 if (self->size == -1) {
935 assert(self->mode != MODE_READ_EOF);
936 for (;;) {
937 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000938 chunksize = BZ2_bzRead(&bzerror, self->fp,
939 buffer, buffersize);
Georg Brandl47fab922006-02-18 21:57:25 +0000940 self->pos += chunksize;
941 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000942
Georg Brandl47fab922006-02-18 21:57:25 +0000943 bytesread += chunksize;
944 if (bzerror == BZ_STREAM_END) {
945 break;
946 } else if (bzerror != BZ_OK) {
947 Util_CatchBZ2Error(bzerror);
948 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000949 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000950 }
Georg Brandl47fab922006-02-18 21:57:25 +0000951 self->mode = MODE_READ_EOF;
952 self->size = self->pos;
953 bytesread = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000954 }
Georg Brandl47fab922006-02-18 21:57:25 +0000955 offset = self->size + offset;
956 } else if (where == 1) {
957 offset = self->pos + offset;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000958 }
959
Guido van Rossum522a6c62007-05-22 23:13:45 +0000960 /* Before getting here, offset must be the absolute position the file
Georg Brandl47fab922006-02-18 21:57:25 +0000961 * pointer should be set to. */
962
963 if (offset >= self->pos) {
964 /* we can move forward */
965 offset -= self->pos;
966 } else {
967 /* we cannot move back, so rewind the stream */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000968 BZ2_bzReadClose(&bzerror, self->fp);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000969 if (bzerror != BZ_OK) {
970 Util_CatchBZ2Error(bzerror);
971 goto cleanup;
972 }
Guido van Rossumf09ca142007-06-13 00:03:05 +0000973 rewind(self->rawfp);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000974 self->pos = 0;
Guido van Rossumf09ca142007-06-13 00:03:05 +0000975 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000976 0, 0, NULL, 0);
977 if (bzerror != BZ_OK) {
978 Util_CatchBZ2Error(bzerror);
979 goto cleanup;
980 }
981 self->mode = MODE_READ;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000982 }
983
Georg Brandl47fab922006-02-18 21:57:25 +0000984 if (offset <= 0 || self->mode == MODE_READ_EOF)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000985 goto exit;
986
987 /* Before getting here, offset must be set to the number of bytes
988 * to walk forward. */
989 for (;;) {
Georg Brandla8bcecc2005-09-03 07:49:53 +0000990 if (offset-bytesread > buffersize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000991 readsize = buffersize;
992 else
Georg Brandla8bcecc2005-09-03 07:49:53 +0000993 /* 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);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000997 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000998 chunksize = BZ2_bzRead(&bzerror, self->fp, buffer, readsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000999 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 }
1013
1014exit:
1015 Py_INCREF(Py_None);
1016 ret = Py_None;
1017
1018cleanup:
1019 RELEASE_LOCK(self);
1020 return ret;
1021}
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{
1032 PyObject *ret = NULL;
1033
1034 if (self->mode == MODE_CLOSED) {
1035 PyErr_SetString(PyExc_ValueError,
1036 "I/O operation on closed file");
1037 goto cleanup;
1038 }
1039
Georg Brandla8bcecc2005-09-03 07:49:53 +00001040#if !defined(HAVE_LARGEFILE_SUPPORT)
Christian Heimes217cfd12007-12-02 14:31:20 +00001041 ret = PyLong_FromLong(self->pos);
Georg Brandla8bcecc2005-09-03 07:49:53 +00001042#else
1043 ret = PyLong_FromLongLong(self->pos);
1044#endif
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001045
1046cleanup:
1047 return ret;
1048}
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{
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001061 PyObject *ret = NULL;
1062 int bzerror = BZ_OK;
1063
Guido van Rossumf09ca142007-06-13 00:03:05 +00001064 if (self->mode == MODE_CLOSED) {
1065 Py_RETURN_NONE;
1066 }
1067
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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;
Guido van Rossumf09ca142007-06-13 00:03:05 +00001080 fclose(self->rawfp);
1081 self->rawfp = NULL;
1082 if (bzerror == BZ_OK) {
1083 Py_INCREF(Py_None);
1084 ret = Py_None;
1085 }
1086 else {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001087 Util_CatchBZ2Error(bzerror);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001088 }
1089
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001090 RELEASE_LOCK(self);
1091 return ret;
1092}
1093
Antoine Pitrou308705e2009-01-10 16:22:51 +00001094PyDoc_STRVAR(BZ2File_enter_doc,
1095"__enter__() -> self.");
1096
1097static PyObject *
1098BZ2File_enter(BZ2FileObject *self)
1099{
1100 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;
1107}
1108
1109PyDoc_STRVAR(BZ2File_exit_doc,
1110"__exit__(*excinfo) -> None. Closes the file.");
1111
1112static PyObject *
1113BZ2File_exit(BZ2FileObject *self, PyObject *args)
1114{
1115 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;
1121}
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[] = {
1127 {"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__},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001134 {"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
Antoine Pitrou308705e2009-01-10 16:22:51 +00001135 {"__enter__", (PyCFunction)BZ2File_enter, METH_NOARGS, BZ2File_enter_doc},
1136 {"__exit__", (PyCFunction)BZ2File_exit, METH_VARARGS, BZ2File_exit_doc},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001137 {NULL, NULL} /* sentinel */
1138};
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{
Christian Heimes217cfd12007-12-02 14:31:20 +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[] = {
1151 {"closed", (getter)BZ2File_get_closed, NULL,
1152 "True if the file is closed"},
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001153 {NULL} /* Sentinel */
1154};
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{
Martin v. Löwis15e62742006-02-27 16:46:16 +00001163 static char *kwlist[] = {"filename", "mode", "buffering",
Guido van Rossumf09ca142007-06-13 00:03:05 +00001164 "compresslevel", 0};
Victor Stinner95503232010-04-23 11:19:05 +00001165 PyObject *name_obj = NULL;
Guido van Rossumf09ca142007-06-13 00:03:05 +00001166 char *name;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001167 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
1173 self->size = -1;
Tim Peterse3228092002-11-09 04:21:44 +00001174
Victor Stinner95503232010-04-23 11:19:05 +00001175 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&|sii:BZ2File",
1176 kwlist, PyUnicode_FSConverter, &name_obj,
1177 &mode, &buffering,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001178 &compresslevel))
1179 return -1;
1180
Victor Stinner95503232010-04-23 11:19:05 +00001181 if (PyBytes_Check(name_obj))
1182 name = PyBytes_AsString(name_obj);
1183 else
1184 name = PyByteArray_AsString(name_obj);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001185 if (compresslevel < 1 || compresslevel > 9) {
1186 PyErr_SetString(PyExc_ValueError,
1187 "compresslevel must be between 1 and 9");
Victor Stinner95503232010-04-23 11:19:05 +00001188 Py_DECREF(name_obj);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001189 return -1;
1190 }
1191
1192 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;
1201
1202 case 'b':
1203 break;
1204
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001205 default:
1206 error = 1;
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001207 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001208 }
1209 if (error) {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001210 PyErr_Format(PyExc_ValueError,
1211 "invalid mode char %c", *mode);
Victor Stinner95503232010-04-23 11:19:05 +00001212 Py_DECREF(name_obj);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001213 return -1;
1214 }
1215 mode++;
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001216 if (*mode == '\0')
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001217 break;
1218 }
1219
Georg Brandl6b95f1d2005-06-03 19:47:00 +00001220 if (mode_char == 0) {
1221 mode_char = 'r';
1222 }
1223
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001224 mode = (mode_char == 'r') ? "rb" : "wb";
Tim Peterse3228092002-11-09 04:21:44 +00001225
Guido van Rossumf09ca142007-06-13 00:03:05 +00001226 self->rawfp = fopen(name, mode);
Victor Stinner95503232010-04-23 11:19:05 +00001227 Py_DECREF(name_obj);
Guido van Rossumf09ca142007-06-13 00:03:05 +00001228 if (self->rawfp == NULL) {
1229 PyErr_SetFromErrno(PyExc_IOError);
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001230 return -1;
Guido van Rossumf09ca142007-06-13 00:03:05 +00001231 }
1232 /* XXX Ignore buffering */
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001233
1234 /* 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
1238 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001239 if (!self->lock) {
1240 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001241 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001242 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001243#endif
1244
1245 if (mode_char == 'r')
Guido van Rossumf09ca142007-06-13 00:03:05 +00001246 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001247 0, 0, NULL, 0);
1248 else
Guido van Rossumf09ca142007-06-13 00:03:05 +00001249 self->fp = BZ2_bzWriteOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001250 compresslevel, 0, 0);
1251
1252 if (bzerror != BZ_OK) {
1253 Util_CatchBZ2Error(bzerror);
1254 goto error;
1255 }
1256
1257 self->mode = (mode_char == 'r') ? MODE_READ : MODE_WRITE;
1258
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001259 return 0;
1260
1261error:
Guido van Rossumf09ca142007-06-13 00:03:05 +00001262 fclose(self->rawfp);
1263 self->rawfp = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001264#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001265 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001266 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001267 self->lock = NULL;
1268 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001269#endif
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001270 return -1;
1271}
1272
1273static void
1274BZ2File_dealloc(BZ2FileObject *self)
1275{
1276 int bzerror;
1277#ifdef WITH_THREAD
1278 if (self->lock)
1279 PyThread_free_lock(self->lock);
1280#endif
1281 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 }
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001291 Util_DropReadAhead(self);
Guido van Rossumf09ca142007-06-13 00:03:05 +00001292 if (self->rawfp != NULL)
1293 fclose(self->rawfp);
Christian Heimes90aa7642007-12-19 02:45:37 +00001294 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{
1301 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;
1308}
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{
Christian Heimes72b710a2008-05-26 13:28:38 +00001315 PyBytesObject* ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001316 ACQUIRE_LOCK(self);
1317 if (self->mode == MODE_CLOSED) {
Georg Brandl86b2fb92008-07-16 03:43:04 +00001318 RELEASE_LOCK(self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001319 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);
Christian Heimes72b710a2008-05-26 13:28:38 +00001325 if (ret == NULL || PyBytes_GET_SIZE(ret) == 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001326 Py_XDECREF(ret);
1327 return NULL;
1328 }
1329 return (PyObject *)ret;
1330}
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 = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001348 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001349 "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*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00001356 0, /*tp_reserved*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001357 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001364 PyObject_GenericGetAttr,/*tp_getattro*/
1365 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001366 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*/
Guido van Rossum79139b22007-02-09 23:20:19 +00001376 0, /*tp_members*/
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001377 BZ2File_getset, /*tp_getset*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001378 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001384 PyType_GenericAlloc, /*tp_alloc*/
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001385 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001386 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001387 0, /*tp_is_gc*/
1388};
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{
Martin v. Löwis423be952008-08-13 15:53:07 +00001406 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001407 char *data;
1408 int datasize;
1409 int bufsize = SMALLCHUNK;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001410 PY_LONG_LONG totalout;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001411 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001412 bz_stream *bzs = &self->bzs;
1413 int bzerror;
1414
Martin v. Löwis423be952008-08-13 15:53:07 +00001415 if (!PyArg_ParseTuple(args, "y*:compress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001416 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001417 data = pdata.buf;
1418 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001419
Martin v. Löwis423be952008-08-13 15:53:07 +00001420 if (datasize == 0) {
1421 PyBuffer_Release(&pdata);
Christian Heimes72b710a2008-05-26 13:28:38 +00001422 return PyBytes_FromStringAndSize("", 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00001423 }
Gustavo Niemeyera6e436e2004-02-14 00:02:45 +00001424
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001425 ACQUIRE_LOCK(self);
1426 if (!self->running) {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001427 PyErr_SetString(PyExc_ValueError,
1428 "this object was already flushed");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001429 goto error;
1430 }
1431
Christian Heimes72b710a2008-05-26 13:28:38 +00001432 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001433 if (!ret)
1434 goto error;
1435
1436 bzs->next_in = data;
1437 bzs->avail_in = datasize;
1438 bzs->next_out = BUF(ret);
1439 bzs->avail_out = bufsize;
1440
1441 totalout = BZS_TOTAL_OUT(bzs);
1442
1443 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 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00001451 if (bzs->avail_in == 0)
1452 break; /* no more input data */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001453 if (bzs->avail_out == 0) {
1454 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001455 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001456 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));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001462 }
1463 }
1464
Christian Heimes72b710a2008-05-26 13:28:38 +00001465 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001466 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1467 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001468
1469 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001470 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001471 return ret;
1472
1473error:
1474 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001475 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001476 Py_XDECREF(ret);
1477 return NULL;
1478}
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{
1490 int bufsize = SMALLCHUNK;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001491 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001492 bz_stream *bzs = &self->bzs;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001493 PY_LONG_LONG totalout;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001494 int bzerror;
1495
1496 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;
1503
Christian Heimes72b710a2008-05-26 13:28:38 +00001504 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001505 if (!ret)
1506 goto error;
1507
1508 bzs->next_out = BUF(ret);
1509 bzs->avail_out = bufsize;
1510
1511 totalout = BZS_TOTAL_OUT(bzs);
1512
1513 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);
Christian Heimes72b710a2008-05-26 13:28:38 +00001525 if (_PyBytes_Resize(&ret, bufsize) < 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001526 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 }
1533
Guido van Rossum522a6c62007-05-22 23:13:45 +00001534 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001535 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001536 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1537 goto error;
1538 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001539
1540 RELEASE_LOCK(self);
1541 return ret;
1542
1543error:
1544 RELEASE_LOCK(self);
1545 Py_XDECREF(ret);
1546 return NULL;
1547}
1548
1549static PyMethodDef BZ2Comp_methods[] = {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001550 {"compress", (PyCFunction)BZ2Comp_compress, METH_VARARGS,
1551 BZ2Comp_compress__doc__},
1552 {"flush", (PyCFunction)BZ2Comp_flush, METH_NOARGS,
1553 BZ2Comp_flush__doc__},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001554 {NULL, NULL} /* sentinel */
1555};
1556
1557
1558/* ===================================================================== */
1559/* Slot definitions for BZ2Comp_Type. */
1560
1561static int
1562BZ2Comp_init(BZ2CompObject *self, PyObject *args, PyObject *kwargs)
1563{
1564 int compresslevel = 9;
1565 int bzerror;
Martin v. Löwis15e62742006-02-27 16:46:16 +00001566 static char *kwlist[] = {"compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001567
1568 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:BZ2Compressor",
1569 kwlist, &compresslevel))
1570 return -1;
1571
1572 if (compresslevel < 1 || compresslevel > 9) {
1573 PyErr_SetString(PyExc_ValueError,
1574 "compresslevel must be between 1 and 9");
1575 goto error;
1576 }
1577
1578#ifdef WITH_THREAD
1579 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001580 if (!self->lock) {
1581 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001582 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001583 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001584#endif
1585
1586 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 }
1592
1593 self->running = 1;
1594
1595 return 0;
1596error:
1597#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001598 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001599 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001600 self->lock = NULL;
1601 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001602#endif
1603 return -1;
1604}
1605
1606static void
1607BZ2Comp_dealloc(BZ2CompObject *self)
1608{
1609#ifdef WITH_THREAD
1610 if (self->lock)
1611 PyThread_free_lock(self->lock);
1612#endif
1613 BZ2_bzCompressEnd(&self->bzs);
Christian Heimes90aa7642007-12-19 02:45:37 +00001614 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 = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001631 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001632 "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*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00001639 0, /*tp_reserved*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001640 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001647 PyObject_GenericGetAttr,/*tp_getattro*/
1648 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001649 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001667 PyType_GenericAlloc, /*tp_alloc*/
1668 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001669 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001670 0, /*tp_is_gc*/
1671};
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[] = {
Guido van Rossum33d26892007-08-05 15:29:28 +00001681 {"unused_data", T_OBJECT, OFF(unused_data), READONLY},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001682 {NULL} /* Sentinel */
1683};
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{
Martin v. Löwis423be952008-08-13 15:53:07 +00001702 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001703 char *data;
1704 int datasize;
1705 int bufsize = SMALLCHUNK;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001706 PY_LONG_LONG totalout;
Neal Norwitz18142c02002-11-05 18:17:32 +00001707 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001708 bz_stream *bzs = &self->bzs;
1709 int bzerror;
1710
Martin v. Löwis423be952008-08-13 15:53:07 +00001711 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001712 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001713 data = pdata.buf;
1714 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001715
1716 ACQUIRE_LOCK(self);
1717 if (!self->running) {
1718 PyErr_SetString(PyExc_EOFError, "end of stream was "
1719 "already found");
1720 goto error;
1721 }
1722
Christian Heimes72b710a2008-05-26 13:28:38 +00001723 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001724 if (!ret)
1725 goto error;
1726
1727 bzs->next_in = data;
1728 bzs->avail_in = datasize;
1729 bzs->next_out = BUF(ret);
1730 bzs->avail_out = bufsize;
1731
1732 totalout = BZS_TOTAL_OUT(bzs);
1733
1734 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 =
Christian Heimes72b710a2008-05-26 13:28:38 +00001742 PyBytes_FromStringAndSize(bzs->next_in,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001743 bzs->avail_in);
1744 }
1745 self->running = 0;
1746 break;
1747 }
1748 if (bzerror != BZ_OK) {
1749 Util_CatchBZ2Error(bzerror);
1750 goto error;
1751 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00001752 if (bzs->avail_in == 0)
1753 break; /* no more input data */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001754 if (bzs->avail_out == 0) {
1755 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001756 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001757 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));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001764 }
1765 }
1766
Guido van Rossum522a6c62007-05-22 23:13:45 +00001767 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001768 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001769 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1770 goto error;
1771 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001772
1773 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001774 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001775 return ret;
1776
1777error:
1778 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001779 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001780 Py_XDECREF(ret);
1781 return NULL;
1782}
1783
1784static PyMethodDef BZ2Decomp_methods[] = {
1785 {"decompress", (PyCFunction)BZ2Decomp_decompress, METH_VARARGS, BZ2Decomp_decompress__doc__},
1786 {NULL, NULL} /* sentinel */
1787};
1788
1789
1790/* ===================================================================== */
1791/* Slot definitions for BZ2Decomp_Type. */
1792
1793static int
1794BZ2Decomp_init(BZ2DecompObject *self, PyObject *args, PyObject *kwargs)
1795{
1796 int bzerror;
1797
1798 if (!PyArg_ParseTuple(args, ":BZ2Decompressor"))
1799 return -1;
1800
1801#ifdef WITH_THREAD
1802 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001803 if (!self->lock) {
1804 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001805 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001806 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001807#endif
1808
Christian Heimes72b710a2008-05-26 13:28:38 +00001809 self->unused_data = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001810 if (!self->unused_data)
1811 goto error;
1812
1813 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 }
1819
1820 self->running = 1;
1821
1822 return 0;
1823
1824error:
1825#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001826 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001827 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001828 self->lock = NULL;
1829 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001830#endif
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001831 Py_CLEAR(self->unused_data);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001832 return -1;
1833}
1834
1835static void
1836BZ2Decomp_dealloc(BZ2DecompObject *self)
1837{
1838#ifdef WITH_THREAD
1839 if (self->lock)
1840 PyThread_free_lock(self->lock);
1841#endif
1842 Py_XDECREF(self->unused_data);
1843 BZ2_bzDecompressEnd(&self->bzs);
Christian Heimes90aa7642007-12-19 02:45:37 +00001844 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 = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001860 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001861 "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*/
Mark Dickinsone94c6792009-02-02 20:36:42 +00001868 0, /*tp_reserved*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001869 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001876 PyObject_GenericGetAttr,/*tp_getattro*/
1877 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001878 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*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001896 PyType_GenericAlloc, /*tp_alloc*/
1897 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001898 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001899 0, /*tp_is_gc*/
1900};
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{
1917 int compresslevel=9;
Martin v. Löwis423be952008-08-13 15:53:07 +00001918 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001919 char *data;
1920 int datasize;
1921 int bufsize;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001922 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001923 bz_stream _bzs;
1924 bz_stream *bzs = &_bzs;
1925 int bzerror;
Martin v. Löwis15e62742006-02-27 16:46:16 +00001926 static char *kwlist[] = {"data", "compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001927
Martin v. Löwis423be952008-08-13 15:53:07 +00001928 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|i",
1929 kwlist, &pdata,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001930 &compresslevel))
1931 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001932 data = pdata.buf;
1933 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001934
1935 if (compresslevel < 1 || compresslevel > 9) {
1936 PyErr_SetString(PyExc_ValueError,
1937 "compresslevel must be between 1 and 9");
Martin v. Löwis423be952008-08-13 15:53:07 +00001938 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001939 return NULL;
1940 }
1941
1942 /* 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;
1945
Christian Heimes72b710a2008-05-26 13:28:38 +00001946 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Martin v. Löwis423be952008-08-13 15:53:07 +00001947 if (!ret) {
1948 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001949 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001950 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001951
1952 memset(bzs, 0, sizeof(bz_stream));
1953
1954 bzs->next_in = data;
1955 bzs->avail_in = datasize;
1956 bzs->next_out = BUF(ret);
1957 bzs->avail_out = bufsize;
1958
1959 bzerror = BZ2_bzCompressInit(bzs, compresslevel, 0, 0);
1960 if (bzerror != BZ_OK) {
1961 Util_CatchBZ2Error(bzerror);
Martin v. Löwis423be952008-08-13 15:53:07 +00001962 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001963 Py_DECREF(ret);
1964 return NULL;
1965 }
Tim Peterse3228092002-11-09 04:21:44 +00001966
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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);
Martin v. Löwis423be952008-08-13 15:53:07 +00001976 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001977 Py_DECREF(ret);
1978 return NULL;
1979 }
1980 if (bzs->avail_out == 0) {
1981 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001982 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001983 BZ2_bzCompressEnd(bzs);
Martin v. Löwis423be952008-08-13 15:53:07 +00001984 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001985 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 }
1991
Guido van Rossum522a6c62007-05-22 23:13:45 +00001992 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001993 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +00001994 ret = NULL;
1995 }
1996 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001997 BZ2_bzCompressEnd(bzs);
1998
Martin v. Löwis423be952008-08-13 15:53:07 +00001999 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002000 return ret;
2001}
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{
Martin v. Löwis423be952008-08-13 15:53:07 +00002013 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002014 char *data;
2015 int datasize;
2016 int bufsize = SMALLCHUNK;
2017 PyObject *ret;
2018 bz_stream _bzs;
2019 bz_stream *bzs = &_bzs;
2020 int bzerror;
2021
Martin v. Löwis423be952008-08-13 15:53:07 +00002022 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002023 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00002024 data = pdata.buf;
2025 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002026
Martin v. Löwis423be952008-08-13 15:53:07 +00002027 if (datasize == 0) {
2028 PyBuffer_Release(&pdata);
Christian Heimes72b710a2008-05-26 13:28:38 +00002029 return PyBytes_FromStringAndSize("", 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00002030 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002031
Christian Heimes72b710a2008-05-26 13:28:38 +00002032 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Martin v. Löwis423be952008-08-13 15:53:07 +00002033 if (!ret) {
2034 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002035 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00002036 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002037
2038 memset(bzs, 0, sizeof(bz_stream));
2039
2040 bzs->next_in = data;
2041 bzs->avail_in = datasize;
2042 bzs->next_out = BUF(ret);
2043 bzs->avail_out = bufsize;
2044
2045 bzerror = BZ2_bzDecompressInit(bzs, 0, 0);
2046 if (bzerror != BZ_OK) {
2047 Util_CatchBZ2Error(bzerror);
2048 Py_DECREF(ret);
Martin v. Löwis423be952008-08-13 15:53:07 +00002049 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002050 return NULL;
2051 }
Tim Peterse3228092002-11-09 04:21:44 +00002052
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +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);
Martin v. Löwis423be952008-08-13 15:53:07 +00002062 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002063 Py_DECREF(ret);
2064 return NULL;
2065 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00002066 if (bzs->avail_in == 0) {
2067 BZ2_bzDecompressEnd(bzs);
2068 PyErr_SetString(PyExc_ValueError,
2069 "couldn't find end of stream");
Martin v. Löwis423be952008-08-13 15:53:07 +00002070 PyBuffer_Release(&pdata);
Guido van Rossumd8faa362007-04-27 19:54:29 +00002071 Py_DECREF(ret);
2072 return NULL;
2073 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002074 if (bzs->avail_out == 0) {
2075 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00002076 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002077 BZ2_bzDecompressEnd(bzs);
Martin v. Löwis423be952008-08-13 15:53:07 +00002078 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002079 return NULL;
2080 }
2081 bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
2082 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002083 }
2084 }
2085
Guido van Rossum522a6c62007-05-22 23:13:45 +00002086 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00002087 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +00002088 ret = NULL;
2089 }
2090 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002091 BZ2_bzDecompressEnd(bzs);
Martin v. Löwis423be952008-08-13 15:53:07 +00002092 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002093
2094 return ret;
2095}
2096
2097static PyMethodDef bz2_methods[] = {
2098 {"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 */
2103};
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 = {
2117 PyModuleDef_HEAD_INIT,
2118 "bz2",
2119 bz2__doc__,
2120 -1,
2121 bz2_methods,
2122 NULL,
2123 NULL,
2124 NULL,
2125 NULL
2126};
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{
2131 PyObject *m;
2132
Christian Heimes90aa7642007-12-19 02:45:37 +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
Martin v. Löwis1a214512008-06-11 05:26:20 +00002137 m = PyModule_Create(&bz2module);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002138 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00002139 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002140
Neal Norwitz53cbdaa2007-08-23 21:42:55 +00002141 PyModule_AddObject(m, "__author__", PyUnicode_FromString(__author__));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002142
2143 Py_INCREF(&BZ2File_Type);
2144 PyModule_AddObject(m, "BZ2File", (PyObject *)&BZ2File_Type);
2145
2146 Py_INCREF(&BZ2Comp_Type);
2147 PyModule_AddObject(m, "BZ2Compressor", (PyObject *)&BZ2Comp_Type);
2148
2149 Py_INCREF(&BZ2Decomp_Type);
2150 PyModule_AddObject(m, "BZ2Decompressor", (PyObject *)&BZ2Decomp_Type);
Martin v. Löwis1a214512008-06-11 05:26:20 +00002151 return m;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002152}