blob: 4bf31ee2f920d0f82810881faa82c6dc50e4fd38 [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
81#define ACQUIRE_LOCK(obj) PyThread_acquire_lock(obj->lock, 1)
82#define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)
83#else
84#define ACQUIRE_LOCK(obj)
85#define RELEASE_LOCK(obj)
86#endif
87
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000088/* Bits in f_newlinetypes */
89#define NEWLINE_UNKNOWN 0 /* No newline seen, yet */
90#define NEWLINE_CR 1 /* \r newline seen */
91#define NEWLINE_LF 2 /* \n newline seen */
92#define NEWLINE_CRLF 4 /* \r\n newline seen */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +000093
94/* ===================================================================== */
95/* Structure definitions. */
96
97typedef struct {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +000098 PyObject_HEAD
Guido van Rossumf09ca142007-06-13 00:03:05 +000099 FILE *rawfp;
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000100
101 char* f_buf; /* Allocated readahead buffer */
102 char* f_bufend; /* Points after last occupied position */
103 char* f_bufptr; /* Current buffer position */
104
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000105 BZFILE *fp;
106 int mode;
Georg Brandla8bcecc2005-09-03 07:49:53 +0000107 Py_off_t pos;
108 Py_off_t size;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000109#ifdef WITH_THREAD
110 PyThread_type_lock lock;
111#endif
112} BZ2FileObject;
113
114typedef struct {
115 PyObject_HEAD
116 bz_stream bzs;
117 int running;
118#ifdef WITH_THREAD
119 PyThread_type_lock lock;
120#endif
121} BZ2CompObject;
122
123typedef struct {
124 PyObject_HEAD
125 bz_stream bzs;
126 int running;
127 PyObject *unused_data;
128#ifdef WITH_THREAD
129 PyThread_type_lock lock;
130#endif
131} BZ2DecompObject;
132
133/* ===================================================================== */
134/* Utility functions. */
135
136static int
137Util_CatchBZ2Error(int bzerror)
138{
139 int ret = 0;
140 switch(bzerror) {
141 case BZ_OK:
142 case BZ_STREAM_END:
143 break;
144
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +0000145#ifdef BZ_CONFIG_ERROR
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000146 case BZ_CONFIG_ERROR:
147 PyErr_SetString(PyExc_SystemError,
148 "the bz2 library was not compiled "
149 "correctly");
150 ret = 1;
151 break;
Gustavo Niemeyer7628f1f2003-04-27 06:25:24 +0000152#endif
Tim Peterse3228092002-11-09 04:21:44 +0000153
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000154 case BZ_PARAM_ERROR:
155 PyErr_SetString(PyExc_ValueError,
156 "the bz2 library has received wrong "
157 "parameters");
158 ret = 1;
159 break;
Tim Peterse3228092002-11-09 04:21:44 +0000160
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000161 case BZ_MEM_ERROR:
162 PyErr_NoMemory();
163 ret = 1;
164 break;
165
166 case BZ_DATA_ERROR:
167 case BZ_DATA_ERROR_MAGIC:
168 PyErr_SetString(PyExc_IOError, "invalid data stream");
169 ret = 1;
170 break;
171
172 case BZ_IO_ERROR:
173 PyErr_SetString(PyExc_IOError, "unknown IO error");
174 ret = 1;
175 break;
176
177 case BZ_UNEXPECTED_EOF:
178 PyErr_SetString(PyExc_EOFError,
179 "compressed file ended before the "
180 "logical end-of-stream was detected");
181 ret = 1;
182 break;
183
184 case BZ_SEQUENCE_ERROR:
185 PyErr_SetString(PyExc_RuntimeError,
186 "wrong sequence of bz2 library "
187 "commands used");
188 ret = 1;
189 break;
190 }
191 return ret;
192}
193
194#if BUFSIZ < 8192
195#define SMALLCHUNK 8192
196#else
197#define SMALLCHUNK BUFSIZ
198#endif
199
200#if SIZEOF_INT < 4
201#define BIGCHUNK (512 * 32)
202#else
203#define BIGCHUNK (512 * 1024)
204#endif
205
206/* This is a hacked version of Python's fileobject.c:new_buffersize(). */
207static size_t
208Util_NewBufferSize(size_t currentsize)
209{
210 if (currentsize > SMALLCHUNK) {
211 /* Keep doubling until we reach BIGCHUNK;
212 then keep adding BIGCHUNK. */
213 if (currentsize <= BIGCHUNK)
214 return currentsize + currentsize;
215 else
216 return currentsize + BIGCHUNK;
217 }
218 return currentsize + SMALLCHUNK;
219}
220
221/* This is a hacked version of Python's fileobject.c:get_line(). */
222static PyObject *
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000223Util_GetLine(BZ2FileObject *f, int n)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000224{
225 char c;
226 char *buf, *end;
227 size_t total_v_size; /* total # of slots in buffer */
228 size_t used_v_size; /* # used slots in buffer */
229 size_t increment; /* amount to increment the buffer */
230 PyObject *v;
231 int bzerror;
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000232 int bytes_read;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000233
234 total_v_size = n > 0 ? n : 100;
Christian Heimes72b710a2008-05-26 13:28:38 +0000235 v = PyBytes_FromStringAndSize((char *)NULL, total_v_size);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000236 if (v == NULL)
237 return NULL;
238
239 buf = BUF(v);
240 end = buf + total_v_size;
241
242 for (;;) {
243 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000244 do {
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000245 bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
Guido van Rossumf09ca142007-06-13 00:03:05 +0000246 f->pos++;
Thomas Wouters1b7f8912007-09-19 03:06:30 +0000247 if (bytes_read == 0)
248 break;
Guido van Rossumf09ca142007-06-13 00:03:05 +0000249 *buf++ = c;
250 } while (bzerror == BZ_OK && c != '\n' && buf != end);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000251 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000252 if (bzerror == BZ_STREAM_END) {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000253 f->size = f->pos;
254 f->mode = MODE_READ_EOF;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000255 break;
256 } else if (bzerror != BZ_OK) {
257 Util_CatchBZ2Error(bzerror);
258 Py_DECREF(v);
259 return NULL;
260 }
261 if (c == '\n')
262 break;
263 /* Must be because buf == end */
264 if (n > 0)
265 break;
266 used_v_size = total_v_size;
267 increment = total_v_size >> 2; /* mild exponential growth */
268 total_v_size += increment;
269 if (total_v_size > INT_MAX) {
270 PyErr_SetString(PyExc_OverflowError,
271 "line is longer than a Python string can hold");
272 Py_DECREF(v);
273 return NULL;
274 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000275 if (_PyBytes_Resize(&v, total_v_size) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000276 return NULL;
Guido van Rossum522a6c62007-05-22 23:13:45 +0000277 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000278 buf = BUF(v) + used_v_size;
279 end = BUF(v) + total_v_size;
280 }
281
282 used_v_size = buf - BUF(v);
Guido van Rossum522a6c62007-05-22 23:13:45 +0000283 if (used_v_size != total_v_size) {
Christian Heimes72b710a2008-05-26 13:28:38 +0000284 if (_PyBytes_Resize(&v, used_v_size) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000285 v = NULL;
286 }
287 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000288 return v;
289}
290
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000291/* This is a hacked version of Python's fileobject.c:drop_readahead(). */
292static void
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000293Util_DropReadAhead(BZ2FileObject *f)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000294{
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000295 if (f->f_buf != NULL) {
296 PyMem_Free(f->f_buf);
297 f->f_buf = NULL;
298 }
299}
300
301/* This is a hacked version of Python's fileobject.c:readahead(). */
302static int
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000303Util_ReadAhead(BZ2FileObject *f, int bufsize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000304{
305 int chunksize;
306 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000307
308 if (f->f_buf != NULL) {
Tim Peterse3228092002-11-09 04:21:44 +0000309 if((f->f_bufend - f->f_bufptr) >= 1)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000310 return 0;
311 else
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000312 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000313 }
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000314 if (f->mode == MODE_READ_EOF) {
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000315 f->f_bufptr = f->f_buf;
316 f->f_bufend = f->f_buf;
317 return 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000318 }
319 if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {
Georg Brandl6aa2d1f2008-08-12 08:35:52 +0000320 PyErr_NoMemory();
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000321 return -1;
322 }
323 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000324 chunksize = BZ2_bzRead(&bzerror, f->fp, f->f_buf, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000325 Py_END_ALLOW_THREADS
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000326 f->pos += chunksize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000327 if (bzerror == BZ_STREAM_END) {
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000328 f->size = f->pos;
329 f->mode = MODE_READ_EOF;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000330 } else if (bzerror != BZ_OK) {
331 Util_CatchBZ2Error(bzerror);
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000332 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000333 return -1;
334 }
335 f->f_bufptr = f->f_buf;
336 f->f_bufend = f->f_buf + chunksize;
337 return 0;
338}
339
340/* This is a hacked version of Python's
341 * fileobject.c:readahead_get_line_skip(). */
Christian Heimes72b710a2008-05-26 13:28:38 +0000342static PyBytesObject *
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000343Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000344{
Christian Heimes72b710a2008-05-26 13:28:38 +0000345 PyBytesObject* s;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000346 char *bufptr;
347 char *buf;
348 int len;
349
350 if (f->f_buf == NULL)
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000351 if (Util_ReadAhead(f, bufsize) < 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000352 return NULL;
353
354 len = f->f_bufend - f->f_bufptr;
Tim Peterse3228092002-11-09 04:21:44 +0000355 if (len == 0)
Christian Heimes72b710a2008-05-26 13:28:38 +0000356 return (PyBytesObject *)
357 PyBytes_FromStringAndSize(NULL, skip);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000358 bufptr = memchr(f->f_bufptr, '\n', len);
359 if (bufptr != NULL) {
360 bufptr++; /* Count the '\n' */
361 len = bufptr - f->f_bufptr;
Christian Heimes72b710a2008-05-26 13:28:38 +0000362 s = (PyBytesObject *)
363 PyBytes_FromStringAndSize(NULL, skip+len);
Tim Peterse3228092002-11-09 04:21:44 +0000364 if (s == NULL)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000365 return NULL;
Christian Heimes72b710a2008-05-26 13:28:38 +0000366 memcpy(PyBytes_AS_STRING(s)+skip, f->f_bufptr, len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000367 f->f_bufptr = bufptr;
368 if (bufptr == f->f_bufend)
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000369 Util_DropReadAhead(f);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000370 } else {
371 bufptr = f->f_bufptr;
372 buf = f->f_buf;
373 f->f_buf = NULL; /* Force new readahead buffer */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +0000374 s = Util_ReadAheadGetLineSkip(f, skip+len,
375 bufsize + (bufsize>>2));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000376 if (s == NULL) {
377 PyMem_Free(buf);
378 return NULL;
379 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000380 memcpy(PyBytes_AS_STRING(s)+skip, bufptr, len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000381 PyMem_Free(buf);
382 }
383 return s;
384}
385
386/* ===================================================================== */
387/* Methods of BZ2File. */
388
389PyDoc_STRVAR(BZ2File_read__doc__,
390"read([size]) -> string\n\
391\n\
392Read at most size uncompressed bytes, returned as a string. If the size\n\
393argument is negative or omitted, read until EOF is reached.\n\
394");
395
396/* This is a hacked version of Python's fileobject.c:file_read(). */
397static PyObject *
398BZ2File_read(BZ2FileObject *self, PyObject *args)
399{
400 long bytesrequested = -1;
401 size_t bytesread, buffersize, chunksize;
402 int bzerror;
403 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000404
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000405 if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))
406 return NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000407
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000408 ACQUIRE_LOCK(self);
409 switch (self->mode) {
410 case MODE_READ:
411 break;
412 case MODE_READ_EOF:
Christian Heimes72b710a2008-05-26 13:28:38 +0000413 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000414 goto cleanup;
415 case MODE_CLOSED:
416 PyErr_SetString(PyExc_ValueError,
417 "I/O operation on closed file");
418 goto cleanup;
419 default:
420 PyErr_SetString(PyExc_IOError,
421 "file is not ready for reading");
422 goto cleanup;
423 }
424
425 if (bytesrequested < 0)
426 buffersize = Util_NewBufferSize((size_t)0);
427 else
428 buffersize = bytesrequested;
429 if (buffersize > INT_MAX) {
430 PyErr_SetString(PyExc_OverflowError,
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +0000431 "requested number of bytes is "
432 "more than a Python string can hold");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000433 goto cleanup;
434 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000435 ret = PyBytes_FromStringAndSize((char *)NULL, buffersize);
Guido van Rossum75c26bc2007-08-07 23:29:20 +0000436 if (ret == NULL || buffersize == 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000437 goto cleanup;
438 bytesread = 0;
439
440 for (;;) {
441 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000442 chunksize = BZ2_bzRead(&bzerror, self->fp,
443 BUF(ret)+bytesread,
444 buffersize-bytesread);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000445 self->pos += chunksize;
446 Py_END_ALLOW_THREADS
447 bytesread += chunksize;
448 if (bzerror == BZ_STREAM_END) {
449 self->size = self->pos;
450 self->mode = MODE_READ_EOF;
451 break;
452 } else if (bzerror != BZ_OK) {
453 Util_CatchBZ2Error(bzerror);
454 Py_DECREF(ret);
455 ret = NULL;
456 goto cleanup;
457 }
458 if (bytesrequested < 0) {
459 buffersize = Util_NewBufferSize(buffersize);
Christian Heimes72b710a2008-05-26 13:28:38 +0000460 if (_PyBytes_Resize(&ret, buffersize) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000461 ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000462 goto cleanup;
Guido van Rossum522a6c62007-05-22 23:13:45 +0000463 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000464 } else {
465 break;
466 }
467 }
Guido van Rossum522a6c62007-05-22 23:13:45 +0000468 if (bytesread != buffersize) {
Christian Heimes72b710a2008-05-26 13:28:38 +0000469 if (_PyBytes_Resize(&ret, bytesread) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +0000470 ret = NULL;
471 }
472 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000473
474cleanup:
475 RELEASE_LOCK(self);
476 return ret;
477}
478
479PyDoc_STRVAR(BZ2File_readline__doc__,
480"readline([size]) -> string\n\
481\n\
482Return the next line from the file, as a string, retaining newline.\n\
483A non-negative size argument will limit the maximum number of bytes to\n\
484return (an incomplete line may be returned then). Return an empty\n\
485string at EOF.\n\
486");
487
488static PyObject *
489BZ2File_readline(BZ2FileObject *self, PyObject *args)
490{
491 PyObject *ret = NULL;
492 int sizehint = -1;
493
494 if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))
495 return NULL;
496
497 ACQUIRE_LOCK(self);
498 switch (self->mode) {
499 case MODE_READ:
500 break;
501 case MODE_READ_EOF:
Christian Heimes72b710a2008-05-26 13:28:38 +0000502 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000503 goto cleanup;
504 case MODE_CLOSED:
505 PyErr_SetString(PyExc_ValueError,
506 "I/O operation on closed file");
507 goto cleanup;
508 default:
509 PyErr_SetString(PyExc_IOError,
510 "file is not ready for reading");
511 goto cleanup;
512 }
513
514 if (sizehint == 0)
Christian Heimes72b710a2008-05-26 13:28:38 +0000515 ret = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000516 else
517 ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);
518
519cleanup:
520 RELEASE_LOCK(self);
521 return ret;
522}
523
524PyDoc_STRVAR(BZ2File_readlines__doc__,
525"readlines([size]) -> list\n\
526\n\
527Call readline() repeatedly and return a list of lines read.\n\
528The optional size argument, if given, is an approximate bound on the\n\
529total number of bytes in the lines returned.\n\
530");
531
532/* This is a hacked version of Python's fileobject.c:file_readlines(). */
533static PyObject *
534BZ2File_readlines(BZ2FileObject *self, PyObject *args)
535{
536 long sizehint = 0;
537 PyObject *list = NULL;
538 PyObject *line;
539 char small_buffer[SMALLCHUNK];
540 char *buffer = small_buffer;
541 size_t buffersize = SMALLCHUNK;
542 PyObject *big_buffer = NULL;
543 size_t nfilled = 0;
544 size_t nread;
545 size_t totalread = 0;
546 char *p, *q, *end;
547 int err;
548 int shortread = 0;
549 int bzerror;
550
551 if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))
552 return NULL;
553
554 ACQUIRE_LOCK(self);
555 switch (self->mode) {
556 case MODE_READ:
557 break;
558 case MODE_READ_EOF:
559 list = PyList_New(0);
560 goto cleanup;
561 case MODE_CLOSED:
562 PyErr_SetString(PyExc_ValueError,
563 "I/O operation on closed file");
564 goto cleanup;
565 default:
566 PyErr_SetString(PyExc_IOError,
567 "file is not ready for reading");
568 goto cleanup;
569 }
570
571 if ((list = PyList_New(0)) == NULL)
572 goto cleanup;
573
574 for (;;) {
575 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000576 nread = BZ2_bzRead(&bzerror, self->fp,
577 buffer+nfilled, buffersize-nfilled);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000578 self->pos += nread;
579 Py_END_ALLOW_THREADS
580 if (bzerror == BZ_STREAM_END) {
581 self->size = self->pos;
582 self->mode = MODE_READ_EOF;
583 if (nread == 0) {
584 sizehint = 0;
585 break;
586 }
587 shortread = 1;
588 } else if (bzerror != BZ_OK) {
589 Util_CatchBZ2Error(bzerror);
590 error:
591 Py_DECREF(list);
592 list = NULL;
593 goto cleanup;
594 }
595 totalread += nread;
596 p = memchr(buffer+nfilled, '\n', nread);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000597 if (!shortread && p == NULL) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000598 /* Need a larger buffer to fit this line */
599 nfilled += nread;
600 buffersize *= 2;
601 if (buffersize > INT_MAX) {
602 PyErr_SetString(PyExc_OverflowError,
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000603 "line is longer than a Python string can hold");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000604 goto error;
605 }
606 if (big_buffer == NULL) {
607 /* Create the big buffer */
Christian Heimes72b710a2008-05-26 13:28:38 +0000608 big_buffer = PyBytes_FromStringAndSize(
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000609 NULL, buffersize);
610 if (big_buffer == NULL)
611 goto error;
Christian Heimes72b710a2008-05-26 13:28:38 +0000612 buffer = PyBytes_AS_STRING(big_buffer);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000613 memcpy(buffer, small_buffer, nfilled);
614 }
615 else {
616 /* Grow the big buffer */
Christian Heimes72b710a2008-05-26 13:28:38 +0000617 if (_PyBytes_Resize(&big_buffer, buffersize) < 0){
Guido van Rossum522a6c62007-05-22 23:13:45 +0000618 big_buffer = NULL;
619 goto error;
620 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000621 buffer = PyBytes_AS_STRING(big_buffer);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000622 }
Guido van Rossum522a6c62007-05-22 23:13:45 +0000623 continue;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000624 }
625 end = buffer+nfilled+nread;
626 q = buffer;
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000627 while (p != NULL) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000628 /* Process complete lines */
629 p++;
Christian Heimes72b710a2008-05-26 13:28:38 +0000630 line = PyBytes_FromStringAndSize(q, p-q);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000631 if (line == NULL)
632 goto error;
633 err = PyList_Append(list, line);
634 Py_DECREF(line);
635 if (err != 0)
636 goto error;
637 q = p;
638 p = memchr(q, '\n', end-q);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000639 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000640 /* Move the remaining incomplete line to the start */
641 nfilled = end-q;
642 memmove(buffer, q, nfilled);
643 if (sizehint > 0)
644 if (totalread >= (size_t)sizehint)
645 break;
646 if (shortread) {
647 sizehint = 0;
648 break;
649 }
650 }
651 if (nfilled != 0) {
652 /* Partial last line */
Christian Heimes72b710a2008-05-26 13:28:38 +0000653 line = PyBytes_FromStringAndSize(buffer, nfilled);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000654 if (line == NULL)
655 goto error;
656 if (sizehint > 0) {
657 /* Need to complete the last line */
658 PyObject *rest = Util_GetLine(self, 0);
659 if (rest == NULL) {
660 Py_DECREF(line);
661 goto error;
662 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000663 PyBytes_Concat(&line, rest);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000664 Py_DECREF(rest);
665 if (line == NULL)
666 goto error;
667 }
668 err = PyList_Append(list, line);
669 Py_DECREF(line);
670 if (err != 0)
671 goto error;
672 }
673
674 cleanup:
675 RELEASE_LOCK(self);
676 if (big_buffer) {
677 Py_DECREF(big_buffer);
678 }
679 return list;
680}
681
682PyDoc_STRVAR(BZ2File_write__doc__,
683"write(data) -> None\n\
684\n\
685Write the 'data' string to file. Note that due to buffering, close() may\n\
686be needed before the file on disk reflects the data written.\n\
687");
688
689/* This is a hacked version of Python's fileobject.c:file_write(). */
690static PyObject *
691BZ2File_write(BZ2FileObject *self, PyObject *args)
692{
693 PyObject *ret = NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +0000694 Py_buffer pbuf;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000695 char *buf;
696 int len;
697 int bzerror;
698
Martin v. Löwis423be952008-08-13 15:53:07 +0000699 if (!PyArg_ParseTuple(args, "y*:write", &pbuf))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000700 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +0000701 buf = pbuf.buf;
702 len = pbuf.len;
Tim Peterse3228092002-11-09 04:21:44 +0000703
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000704 ACQUIRE_LOCK(self);
705 switch (self->mode) {
706 case MODE_WRITE:
707 break;
Tim Peterse3228092002-11-09 04:21:44 +0000708
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000709 case MODE_CLOSED:
710 PyErr_SetString(PyExc_ValueError,
711 "I/O operation on closed file");
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000712 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000713
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000714 default:
715 PyErr_SetString(PyExc_IOError,
716 "file is not ready for writing");
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000717 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000718 }
719
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000720 Py_BEGIN_ALLOW_THREADS
721 BZ2_bzWrite (&bzerror, self->fp, buf, len);
722 self->pos += len;
723 Py_END_ALLOW_THREADS
Tim Peterse3228092002-11-09 04:21:44 +0000724
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000725 if (bzerror != BZ_OK) {
726 Util_CatchBZ2Error(bzerror);
727 goto cleanup;
728 }
Tim Peterse3228092002-11-09 04:21:44 +0000729
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000730 Py_INCREF(Py_None);
731 ret = Py_None;
732
733cleanup:
Martin v. Löwis423be952008-08-13 15:53:07 +0000734 PyBuffer_Release(&pbuf);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000735 RELEASE_LOCK(self);
736 return ret;
737}
738
739PyDoc_STRVAR(BZ2File_writelines__doc__,
740"writelines(sequence_of_strings) -> None\n\
741\n\
742Write the sequence of strings to the file. Note that newlines are not\n\
743added. The sequence can be any iterable object producing strings. This is\n\
744equivalent to calling write() for each string.\n\
745");
746
747/* This is a hacked version of Python's fileobject.c:file_writelines(). */
748static PyObject *
749BZ2File_writelines(BZ2FileObject *self, PyObject *seq)
750{
751#define CHUNKSIZE 1000
752 PyObject *list = NULL;
753 PyObject *iter = NULL;
754 PyObject *ret = NULL;
755 PyObject *line;
756 int i, j, index, len, islist;
757 int bzerror;
758
759 ACQUIRE_LOCK(self);
Thomas Wouters00ee7ba2006-08-21 19:07:27 +0000760 switch (self->mode) {
761 case MODE_WRITE:
762 break;
763
764 case MODE_CLOSED:
765 PyErr_SetString(PyExc_ValueError,
766 "I/O operation on closed file");
767 goto error;
768
769 default:
770 PyErr_SetString(PyExc_IOError,
771 "file is not ready for writing");
772 goto error;
773 }
774
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000775 islist = PyList_Check(seq);
776 if (!islist) {
777 iter = PyObject_GetIter(seq);
778 if (iter == NULL) {
779 PyErr_SetString(PyExc_TypeError,
780 "writelines() requires an iterable argument");
781 goto error;
782 }
783 list = PyList_New(CHUNKSIZE);
784 if (list == NULL)
785 goto error;
786 }
787
788 /* Strategy: slurp CHUNKSIZE lines into a private list,
789 checking that they are all strings, then write that list
790 without holding the interpreter lock, then come back for more. */
791 for (index = 0; ; index += CHUNKSIZE) {
792 if (islist) {
793 Py_XDECREF(list);
794 list = PyList_GetSlice(seq, index, index+CHUNKSIZE);
795 if (list == NULL)
796 goto error;
797 j = PyList_GET_SIZE(list);
798 }
799 else {
800 for (j = 0; j < CHUNKSIZE; j++) {
801 line = PyIter_Next(iter);
802 if (line == NULL) {
803 if (PyErr_Occurred())
804 goto error;
805 break;
806 }
807 PyList_SetItem(list, j, line);
808 }
809 }
810 if (j == 0)
811 break;
812
Guido van Rossum522a6c62007-05-22 23:13:45 +0000813 /* Check that all entries are indeed byte strings. If not,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000814 apply the same rules as for file.write() and
815 convert the rets to strings. This is slow, but
816 seems to be the only way since all conversion APIs
817 could potentially execute Python code. */
818 for (i = 0; i < j; i++) {
819 PyObject *v = PyList_GET_ITEM(list, i);
Christian Heimes72b710a2008-05-26 13:28:38 +0000820 if (!PyBytes_Check(v)) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000821 const char *buffer;
Martin v. Löwis18e16552006-02-15 17:27:45 +0000822 Py_ssize_t len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000823 if (PyObject_AsCharBuffer(v, &buffer, &len)) {
824 PyErr_SetString(PyExc_TypeError,
825 "writelines() "
826 "argument must be "
827 "a sequence of "
Guido van Rossum522a6c62007-05-22 23:13:45 +0000828 "bytes objects");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000829 goto error;
830 }
Christian Heimes72b710a2008-05-26 13:28:38 +0000831 line = PyBytes_FromStringAndSize(buffer,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000832 len);
833 if (line == NULL)
834 goto error;
835 Py_DECREF(v);
836 PyList_SET_ITEM(list, i, line);
837 }
838 }
839
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000840 /* Since we are releasing the global lock, the
841 following code may *not* execute Python code. */
842 Py_BEGIN_ALLOW_THREADS
843 for (i = 0; i < j; i++) {
844 line = PyList_GET_ITEM(list, i);
Christian Heimes72b710a2008-05-26 13:28:38 +0000845 len = PyBytes_GET_SIZE(line);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000846 BZ2_bzWrite (&bzerror, self->fp,
Christian Heimes72b710a2008-05-26 13:28:38 +0000847 PyBytes_AS_STRING(line), len);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000848 if (bzerror != BZ_OK) {
849 Py_BLOCK_THREADS
850 Util_CatchBZ2Error(bzerror);
851 goto error;
852 }
853 }
854 Py_END_ALLOW_THREADS
855
856 if (j < CHUNKSIZE)
857 break;
858 }
859
860 Py_INCREF(Py_None);
861 ret = Py_None;
862
863 error:
864 RELEASE_LOCK(self);
865 Py_XDECREF(list);
866 Py_XDECREF(iter);
867 return ret;
868#undef CHUNKSIZE
869}
870
871PyDoc_STRVAR(BZ2File_seek__doc__,
872"seek(offset [, whence]) -> None\n\
873\n\
874Move to new file position. Argument offset is a byte count. Optional\n\
875argument whence defaults to 0 (offset from start of file, offset\n\
876should be >= 0); other values are 1 (move relative to current position,\n\
877positive or negative), and 2 (move relative to end of file, usually\n\
878negative, although many platforms allow seeking beyond the end of a file).\n\
879\n\
880Note that seeking of bz2 files is emulated, and depending on the parameters\n\
881the operation may be extremely slow.\n\
882");
883
884static PyObject *
885BZ2File_seek(BZ2FileObject *self, PyObject *args)
886{
887 int where = 0;
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000888 PyObject *offobj;
889 Py_off_t offset;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000890 char small_buffer[SMALLCHUNK];
891 char *buffer = small_buffer;
892 size_t buffersize = SMALLCHUNK;
Thomas Wouters902d6eb2007-01-09 23:18:33 +0000893 Py_off_t bytesread = 0;
Georg Brandla8bcecc2005-09-03 07:49:53 +0000894 size_t readsize;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000895 int chunksize;
896 int bzerror;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000897 PyObject *ret = NULL;
Tim Peterse3228092002-11-09 04:21:44 +0000898
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000899 if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))
900 return NULL;
901#if !defined(HAVE_LARGEFILE_SUPPORT)
Christian Heimes217cfd12007-12-02 14:31:20 +0000902 offset = PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000903#else
904 offset = PyLong_Check(offobj) ?
Christian Heimes217cfd12007-12-02 14:31:20 +0000905 PyLong_AsLongLong(offobj) : PyLong_AsLong(offobj);
Georg Brandl33a5f2a2005-08-21 14:16:04 +0000906#endif
907 if (PyErr_Occurred())
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000908 return NULL;
909
910 ACQUIRE_LOCK(self);
911 Util_DropReadAhead(self);
912 switch (self->mode) {
913 case MODE_READ:
914 case MODE_READ_EOF:
915 break;
Tim Peterse3228092002-11-09 04:21:44 +0000916
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000917 case MODE_CLOSED:
918 PyErr_SetString(PyExc_ValueError,
919 "I/O operation on closed file");
Thomas Wouters89f507f2006-12-13 04:49:30 +0000920 goto cleanup;
Tim Peterse3228092002-11-09 04:21:44 +0000921
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000922 default:
923 PyErr_SetString(PyExc_IOError,
924 "seek works only while reading");
Thomas Wouters89f507f2006-12-13 04:49:30 +0000925 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000926 }
927
Georg Brandl47fab922006-02-18 21:57:25 +0000928 if (where == 2) {
929 if (self->size == -1) {
930 assert(self->mode != MODE_READ_EOF);
931 for (;;) {
932 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000933 chunksize = BZ2_bzRead(&bzerror, self->fp,
934 buffer, buffersize);
Georg Brandl47fab922006-02-18 21:57:25 +0000935 self->pos += chunksize;
936 Py_END_ALLOW_THREADS
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000937
Georg Brandl47fab922006-02-18 21:57:25 +0000938 bytesread += chunksize;
939 if (bzerror == BZ_STREAM_END) {
940 break;
941 } else if (bzerror != BZ_OK) {
942 Util_CatchBZ2Error(bzerror);
943 goto cleanup;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000944 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000945 }
Georg Brandl47fab922006-02-18 21:57:25 +0000946 self->mode = MODE_READ_EOF;
947 self->size = self->pos;
948 bytesread = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000949 }
Georg Brandl47fab922006-02-18 21:57:25 +0000950 offset = self->size + offset;
951 } else if (where == 1) {
952 offset = self->pos + offset;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000953 }
954
Guido van Rossum522a6c62007-05-22 23:13:45 +0000955 /* Before getting here, offset must be the absolute position the file
Georg Brandl47fab922006-02-18 21:57:25 +0000956 * pointer should be set to. */
957
958 if (offset >= self->pos) {
959 /* we can move forward */
960 offset -= self->pos;
961 } else {
962 /* we cannot move back, so rewind the stream */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000963 BZ2_bzReadClose(&bzerror, self->fp);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000964 if (bzerror != BZ_OK) {
965 Util_CatchBZ2Error(bzerror);
966 goto cleanup;
967 }
Guido van Rossumf09ca142007-06-13 00:03:05 +0000968 rewind(self->rawfp);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000969 self->pos = 0;
Guido van Rossumf09ca142007-06-13 00:03:05 +0000970 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000971 0, 0, NULL, 0);
972 if (bzerror != BZ_OK) {
973 Util_CatchBZ2Error(bzerror);
974 goto cleanup;
975 }
976 self->mode = MODE_READ;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000977 }
978
Georg Brandl47fab922006-02-18 21:57:25 +0000979 if (offset <= 0 || self->mode == MODE_READ_EOF)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000980 goto exit;
981
982 /* Before getting here, offset must be set to the number of bytes
983 * to walk forward. */
984 for (;;) {
Georg Brandla8bcecc2005-09-03 07:49:53 +0000985 if (offset-bytesread > buffersize)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000986 readsize = buffersize;
987 else
Georg Brandla8bcecc2005-09-03 07:49:53 +0000988 /* offset might be wider that readsize, but the result
989 * of the subtraction is bound by buffersize (see the
990 * condition above). buffersize is 8192. */
991 readsize = (size_t)(offset-bytesread);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000992 Py_BEGIN_ALLOW_THREADS
Guido van Rossumf09ca142007-06-13 00:03:05 +0000993 chunksize = BZ2_bzRead(&bzerror, self->fp, buffer, readsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +0000994 self->pos += chunksize;
995 Py_END_ALLOW_THREADS
996 bytesread += chunksize;
997 if (bzerror == BZ_STREAM_END) {
998 self->size = self->pos;
999 self->mode = MODE_READ_EOF;
1000 break;
1001 } else if (bzerror != BZ_OK) {
1002 Util_CatchBZ2Error(bzerror);
1003 goto cleanup;
1004 }
1005 if (bytesread == offset)
1006 break;
1007 }
1008
1009exit:
1010 Py_INCREF(Py_None);
1011 ret = Py_None;
1012
1013cleanup:
1014 RELEASE_LOCK(self);
1015 return ret;
1016}
1017
1018PyDoc_STRVAR(BZ2File_tell__doc__,
1019"tell() -> int\n\
1020\n\
1021Return the current file position, an integer (may be a long integer).\n\
1022");
1023
1024static PyObject *
1025BZ2File_tell(BZ2FileObject *self, PyObject *args)
1026{
1027 PyObject *ret = NULL;
1028
1029 if (self->mode == MODE_CLOSED) {
1030 PyErr_SetString(PyExc_ValueError,
1031 "I/O operation on closed file");
1032 goto cleanup;
1033 }
1034
Georg Brandla8bcecc2005-09-03 07:49:53 +00001035#if !defined(HAVE_LARGEFILE_SUPPORT)
Christian Heimes217cfd12007-12-02 14:31:20 +00001036 ret = PyLong_FromLong(self->pos);
Georg Brandla8bcecc2005-09-03 07:49:53 +00001037#else
1038 ret = PyLong_FromLongLong(self->pos);
1039#endif
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001040
1041cleanup:
1042 return ret;
1043}
1044
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001045PyDoc_STRVAR(BZ2File_close__doc__,
1046"close() -> None or (perhaps) an integer\n\
1047\n\
1048Close the file. Sets data attribute .closed to true. A closed file\n\
1049cannot be used for further I/O operations. close() may be called more\n\
1050than once without error.\n\
1051");
1052
1053static PyObject *
1054BZ2File_close(BZ2FileObject *self)
1055{
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001056 PyObject *ret = NULL;
1057 int bzerror = BZ_OK;
1058
Guido van Rossumf09ca142007-06-13 00:03:05 +00001059 if (self->mode == MODE_CLOSED) {
1060 Py_RETURN_NONE;
1061 }
1062
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001063 ACQUIRE_LOCK(self);
1064 switch (self->mode) {
1065 case MODE_READ:
1066 case MODE_READ_EOF:
1067 BZ2_bzReadClose(&bzerror, self->fp);
1068 break;
1069 case MODE_WRITE:
1070 BZ2_bzWriteClose(&bzerror, self->fp,
1071 0, NULL, NULL);
1072 break;
1073 }
1074 self->mode = MODE_CLOSED;
Guido van Rossumf09ca142007-06-13 00:03:05 +00001075 fclose(self->rawfp);
1076 self->rawfp = NULL;
1077 if (bzerror == BZ_OK) {
1078 Py_INCREF(Py_None);
1079 ret = Py_None;
1080 }
1081 else {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001082 Util_CatchBZ2Error(bzerror);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001083 }
1084
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001085 RELEASE_LOCK(self);
1086 return ret;
1087}
1088
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001089static PyObject *BZ2File_getiter(BZ2FileObject *self);
1090
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001091static PyMethodDef BZ2File_methods[] = {
1092 {"read", (PyCFunction)BZ2File_read, METH_VARARGS, BZ2File_read__doc__},
1093 {"readline", (PyCFunction)BZ2File_readline, METH_VARARGS, BZ2File_readline__doc__},
1094 {"readlines", (PyCFunction)BZ2File_readlines, METH_VARARGS, BZ2File_readlines__doc__},
1095 {"write", (PyCFunction)BZ2File_write, METH_VARARGS, BZ2File_write__doc__},
1096 {"writelines", (PyCFunction)BZ2File_writelines, METH_O, BZ2File_writelines__doc__},
1097 {"seek", (PyCFunction)BZ2File_seek, METH_VARARGS, BZ2File_seek__doc__},
1098 {"tell", (PyCFunction)BZ2File_tell, METH_NOARGS, BZ2File_tell__doc__},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001099 {"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
1100 {NULL, NULL} /* sentinel */
1101};
1102
1103
1104/* ===================================================================== */
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001105/* Getters and setters of BZ2File. */
1106
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001107static PyObject *
1108BZ2File_get_closed(BZ2FileObject *self, void *closure)
1109{
Christian Heimes217cfd12007-12-02 14:31:20 +00001110 return PyLong_FromLong(self->mode == MODE_CLOSED);
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001111}
1112
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001113static PyGetSetDef BZ2File_getset[] = {
1114 {"closed", (getter)BZ2File_get_closed, NULL,
1115 "True if the file is closed"},
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001116 {NULL} /* Sentinel */
1117};
1118
1119
1120/* ===================================================================== */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001121/* Slot definitions for BZ2File_Type. */
1122
1123static int
1124BZ2File_init(BZ2FileObject *self, PyObject *args, PyObject *kwargs)
1125{
Martin v. Löwis15e62742006-02-27 16:46:16 +00001126 static char *kwlist[] = {"filename", "mode", "buffering",
Guido van Rossumf09ca142007-06-13 00:03:05 +00001127 "compresslevel", 0};
1128 char *name;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001129 char *mode = "r";
1130 int buffering = -1;
1131 int compresslevel = 9;
1132 int bzerror;
1133 int mode_char = 0;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001134
1135 self->size = -1;
Tim Peterse3228092002-11-09 04:21:44 +00001136
Guido van Rossumf09ca142007-06-13 00:03:05 +00001137 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s|sii:BZ2File",
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001138 kwlist, &name, &mode, &buffering,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001139 &compresslevel))
1140 return -1;
1141
1142 if (compresslevel < 1 || compresslevel > 9) {
1143 PyErr_SetString(PyExc_ValueError,
1144 "compresslevel must be between 1 and 9");
1145 return -1;
1146 }
1147
1148 for (;;) {
1149 int error = 0;
1150 switch (*mode) {
1151 case 'r':
1152 case 'w':
1153 if (mode_char)
1154 error = 1;
1155 mode_char = *mode;
1156 break;
1157
1158 case 'b':
1159 break;
1160
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001161 default:
1162 error = 1;
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001163 break;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001164 }
1165 if (error) {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001166 PyErr_Format(PyExc_ValueError,
1167 "invalid mode char %c", *mode);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001168 return -1;
1169 }
1170 mode++;
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001171 if (*mode == '\0')
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001172 break;
1173 }
1174
Georg Brandl6b95f1d2005-06-03 19:47:00 +00001175 if (mode_char == 0) {
1176 mode_char = 'r';
1177 }
1178
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001179 mode = (mode_char == 'r') ? "rb" : "wb";
Tim Peterse3228092002-11-09 04:21:44 +00001180
Guido van Rossumf09ca142007-06-13 00:03:05 +00001181 self->rawfp = fopen(name, mode);
1182 if (self->rawfp == NULL) {
1183 PyErr_SetFromErrno(PyExc_IOError);
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001184 return -1;
Guido van Rossumf09ca142007-06-13 00:03:05 +00001185 }
1186 /* XXX Ignore buffering */
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001187
1188 /* From now on, we have stuff to dealloc, so jump to error label
1189 * instead of returning */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001190
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001191#ifdef WITH_THREAD
1192 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001193 if (!self->lock) {
1194 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001195 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001196 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001197#endif
1198
1199 if (mode_char == 'r')
Guido van Rossumf09ca142007-06-13 00:03:05 +00001200 self->fp = BZ2_bzReadOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001201 0, 0, NULL, 0);
1202 else
Guido van Rossumf09ca142007-06-13 00:03:05 +00001203 self->fp = BZ2_bzWriteOpen(&bzerror, self->rawfp,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001204 compresslevel, 0, 0);
1205
1206 if (bzerror != BZ_OK) {
1207 Util_CatchBZ2Error(bzerror);
1208 goto error;
1209 }
1210
1211 self->mode = (mode_char == 'r') ? MODE_READ : MODE_WRITE;
1212
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001213 return 0;
1214
1215error:
Guido van Rossumf09ca142007-06-13 00:03:05 +00001216 fclose(self->rawfp);
1217 self->rawfp = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001218#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001219 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001220 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001221 self->lock = NULL;
1222 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001223#endif
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001224 return -1;
1225}
1226
1227static void
1228BZ2File_dealloc(BZ2FileObject *self)
1229{
1230 int bzerror;
1231#ifdef WITH_THREAD
1232 if (self->lock)
1233 PyThread_free_lock(self->lock);
1234#endif
1235 switch (self->mode) {
1236 case MODE_READ:
1237 case MODE_READ_EOF:
1238 BZ2_bzReadClose(&bzerror, self->fp);
1239 break;
1240 case MODE_WRITE:
1241 BZ2_bzWriteClose(&bzerror, self->fp,
1242 0, NULL, NULL);
1243 break;
1244 }
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001245 Util_DropReadAhead(self);
Guido van Rossumf09ca142007-06-13 00:03:05 +00001246 if (self->rawfp != NULL)
1247 fclose(self->rawfp);
Christian Heimes90aa7642007-12-19 02:45:37 +00001248 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001249}
1250
1251/* This is a hacked version of Python's fileobject.c:file_getiter(). */
1252static PyObject *
1253BZ2File_getiter(BZ2FileObject *self)
1254{
1255 if (self->mode == MODE_CLOSED) {
1256 PyErr_SetString(PyExc_ValueError,
1257 "I/O operation on closed file");
1258 return NULL;
1259 }
1260 Py_INCREF((PyObject*)self);
1261 return (PyObject *)self;
1262}
1263
1264/* This is a hacked version of Python's fileobject.c:file_iternext(). */
1265#define READAHEAD_BUFSIZE 8192
1266static PyObject *
1267BZ2File_iternext(BZ2FileObject *self)
1268{
Christian Heimes72b710a2008-05-26 13:28:38 +00001269 PyBytesObject* ret;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001270 ACQUIRE_LOCK(self);
1271 if (self->mode == MODE_CLOSED) {
Georg Brandl86b2fb92008-07-16 03:43:04 +00001272 RELEASE_LOCK(self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001273 PyErr_SetString(PyExc_ValueError,
1274 "I/O operation on closed file");
1275 return NULL;
1276 }
1277 ret = Util_ReadAheadGetLineSkip(self, 0, READAHEAD_BUFSIZE);
1278 RELEASE_LOCK(self);
Christian Heimes72b710a2008-05-26 13:28:38 +00001279 if (ret == NULL || PyBytes_GET_SIZE(ret) == 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001280 Py_XDECREF(ret);
1281 return NULL;
1282 }
1283 return (PyObject *)ret;
1284}
1285
1286/* ===================================================================== */
1287/* BZ2File_Type definition. */
1288
1289PyDoc_VAR(BZ2File__doc__) =
1290PyDoc_STR(
1291"BZ2File(name [, mode='r', buffering=0, compresslevel=9]) -> file object\n\
1292\n\
1293Open a bz2 file. The mode can be 'r' or 'w', for reading (default) or\n\
1294writing. When opened for writing, the file will be created if it doesn't\n\
1295exist, and truncated otherwise. If the buffering argument is given, 0 means\n\
1296unbuffered, and larger numbers specify the buffer size. If compresslevel\n\
1297is given, must be a number between 1 and 9.\n\
Guido van Rossum88e860c2007-06-13 01:46:31 +00001298Data read is always returned in bytes; data written ought to be bytes.\n\
1299");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001300
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001301static PyTypeObject BZ2File_Type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001302 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001303 "bz2.BZ2File", /*tp_name*/
1304 sizeof(BZ2FileObject), /*tp_basicsize*/
1305 0, /*tp_itemsize*/
1306 (destructor)BZ2File_dealloc, /*tp_dealloc*/
1307 0, /*tp_print*/
1308 0, /*tp_getattr*/
1309 0, /*tp_setattr*/
1310 0, /*tp_compare*/
1311 0, /*tp_repr*/
1312 0, /*tp_as_number*/
1313 0, /*tp_as_sequence*/
1314 0, /*tp_as_mapping*/
1315 0, /*tp_hash*/
1316 0, /*tp_call*/
1317 0, /*tp_str*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001318 PyObject_GenericGetAttr,/*tp_getattro*/
1319 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001320 0, /*tp_as_buffer*/
1321 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1322 BZ2File__doc__, /*tp_doc*/
1323 0, /*tp_traverse*/
1324 0, /*tp_clear*/
1325 0, /*tp_richcompare*/
1326 0, /*tp_weaklistoffset*/
1327 (getiterfunc)BZ2File_getiter, /*tp_iter*/
1328 (iternextfunc)BZ2File_iternext, /*tp_iternext*/
1329 BZ2File_methods, /*tp_methods*/
Guido van Rossum79139b22007-02-09 23:20:19 +00001330 0, /*tp_members*/
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001331 BZ2File_getset, /*tp_getset*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001332 0, /*tp_base*/
1333 0, /*tp_dict*/
1334 0, /*tp_descr_get*/
1335 0, /*tp_descr_set*/
1336 0, /*tp_dictoffset*/
1337 (initproc)BZ2File_init, /*tp_init*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001338 PyType_GenericAlloc, /*tp_alloc*/
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001339 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001340 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001341 0, /*tp_is_gc*/
1342};
1343
1344
1345/* ===================================================================== */
1346/* Methods of BZ2Comp. */
1347
1348PyDoc_STRVAR(BZ2Comp_compress__doc__,
1349"compress(data) -> string\n\
1350\n\
1351Provide more data to the compressor object. It will return chunks of\n\
1352compressed data whenever possible. When you've finished providing data\n\
1353to compress, call the flush() method to finish the compression process,\n\
1354and return what is left in the internal buffers.\n\
1355");
1356
1357static PyObject *
1358BZ2Comp_compress(BZ2CompObject *self, PyObject *args)
1359{
Martin v. Löwis423be952008-08-13 15:53:07 +00001360 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001361 char *data;
1362 int datasize;
1363 int bufsize = SMALLCHUNK;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001364 PY_LONG_LONG totalout;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001365 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001366 bz_stream *bzs = &self->bzs;
1367 int bzerror;
1368
Martin v. Löwis423be952008-08-13 15:53:07 +00001369 if (!PyArg_ParseTuple(args, "y*:compress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001370 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001371 data = pdata.buf;
1372 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001373
Martin v. Löwis423be952008-08-13 15:53:07 +00001374 if (datasize == 0) {
1375 PyBuffer_Release(&pdata);
Christian Heimes72b710a2008-05-26 13:28:38 +00001376 return PyBytes_FromStringAndSize("", 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00001377 }
Gustavo Niemeyera6e436e2004-02-14 00:02:45 +00001378
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001379 ACQUIRE_LOCK(self);
1380 if (!self->running) {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001381 PyErr_SetString(PyExc_ValueError,
1382 "this object was already flushed");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001383 goto error;
1384 }
1385
Christian Heimes72b710a2008-05-26 13:28:38 +00001386 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001387 if (!ret)
1388 goto error;
1389
1390 bzs->next_in = data;
1391 bzs->avail_in = datasize;
1392 bzs->next_out = BUF(ret);
1393 bzs->avail_out = bufsize;
1394
1395 totalout = BZS_TOTAL_OUT(bzs);
1396
1397 for (;;) {
1398 Py_BEGIN_ALLOW_THREADS
1399 bzerror = BZ2_bzCompress(bzs, BZ_RUN);
1400 Py_END_ALLOW_THREADS
1401 if (bzerror != BZ_RUN_OK) {
1402 Util_CatchBZ2Error(bzerror);
1403 goto error;
1404 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00001405 if (bzs->avail_in == 0)
1406 break; /* no more input data */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001407 if (bzs->avail_out == 0) {
1408 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001409 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001410 BZ2_bzCompressEnd(bzs);
1411 goto error;
1412 }
1413 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1414 - totalout);
1415 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001416 }
1417 }
1418
Christian Heimes72b710a2008-05-26 13:28:38 +00001419 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001420 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1421 goto error;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001422
1423 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001424 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001425 return ret;
1426
1427error:
1428 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001429 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001430 Py_XDECREF(ret);
1431 return NULL;
1432}
1433
1434PyDoc_STRVAR(BZ2Comp_flush__doc__,
1435"flush() -> string\n\
1436\n\
1437Finish the compression process and return what is left in internal buffers.\n\
1438You must not use the compressor object after calling this method.\n\
1439");
1440
1441static PyObject *
1442BZ2Comp_flush(BZ2CompObject *self)
1443{
1444 int bufsize = SMALLCHUNK;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001445 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001446 bz_stream *bzs = &self->bzs;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001447 PY_LONG_LONG totalout;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001448 int bzerror;
1449
1450 ACQUIRE_LOCK(self);
1451 if (!self->running) {
1452 PyErr_SetString(PyExc_ValueError, "object was already "
1453 "flushed");
1454 goto error;
1455 }
1456 self->running = 0;
1457
Christian Heimes72b710a2008-05-26 13:28:38 +00001458 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001459 if (!ret)
1460 goto error;
1461
1462 bzs->next_out = BUF(ret);
1463 bzs->avail_out = bufsize;
1464
1465 totalout = BZS_TOTAL_OUT(bzs);
1466
1467 for (;;) {
1468 Py_BEGIN_ALLOW_THREADS
1469 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1470 Py_END_ALLOW_THREADS
1471 if (bzerror == BZ_STREAM_END) {
1472 break;
1473 } else if (bzerror != BZ_FINISH_OK) {
1474 Util_CatchBZ2Error(bzerror);
1475 goto error;
1476 }
1477 if (bzs->avail_out == 0) {
1478 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001479 if (_PyBytes_Resize(&ret, bufsize) < 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001480 goto error;
1481 bzs->next_out = BUF(ret);
1482 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1483 - totalout);
1484 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1485 }
1486 }
1487
Guido van Rossum522a6c62007-05-22 23:13:45 +00001488 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001489 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001490 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1491 goto error;
1492 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001493
1494 RELEASE_LOCK(self);
1495 return ret;
1496
1497error:
1498 RELEASE_LOCK(self);
1499 Py_XDECREF(ret);
1500 return NULL;
1501}
1502
1503static PyMethodDef BZ2Comp_methods[] = {
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001504 {"compress", (PyCFunction)BZ2Comp_compress, METH_VARARGS,
1505 BZ2Comp_compress__doc__},
1506 {"flush", (PyCFunction)BZ2Comp_flush, METH_NOARGS,
1507 BZ2Comp_flush__doc__},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001508 {NULL, NULL} /* sentinel */
1509};
1510
1511
1512/* ===================================================================== */
1513/* Slot definitions for BZ2Comp_Type. */
1514
1515static int
1516BZ2Comp_init(BZ2CompObject *self, PyObject *args, PyObject *kwargs)
1517{
1518 int compresslevel = 9;
1519 int bzerror;
Martin v. Löwis15e62742006-02-27 16:46:16 +00001520 static char *kwlist[] = {"compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001521
1522 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:BZ2Compressor",
1523 kwlist, &compresslevel))
1524 return -1;
1525
1526 if (compresslevel < 1 || compresslevel > 9) {
1527 PyErr_SetString(PyExc_ValueError,
1528 "compresslevel must be between 1 and 9");
1529 goto error;
1530 }
1531
1532#ifdef WITH_THREAD
1533 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001534 if (!self->lock) {
1535 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001536 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001537 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001538#endif
1539
1540 memset(&self->bzs, 0, sizeof(bz_stream));
1541 bzerror = BZ2_bzCompressInit(&self->bzs, compresslevel, 0, 0);
1542 if (bzerror != BZ_OK) {
1543 Util_CatchBZ2Error(bzerror);
1544 goto error;
1545 }
1546
1547 self->running = 1;
1548
1549 return 0;
1550error:
1551#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001552 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001553 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001554 self->lock = NULL;
1555 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001556#endif
1557 return -1;
1558}
1559
1560static void
1561BZ2Comp_dealloc(BZ2CompObject *self)
1562{
1563#ifdef WITH_THREAD
1564 if (self->lock)
1565 PyThread_free_lock(self->lock);
1566#endif
1567 BZ2_bzCompressEnd(&self->bzs);
Christian Heimes90aa7642007-12-19 02:45:37 +00001568 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001569}
1570
1571
1572/* ===================================================================== */
1573/* BZ2Comp_Type definition. */
1574
1575PyDoc_STRVAR(BZ2Comp__doc__,
1576"BZ2Compressor([compresslevel=9]) -> compressor object\n\
1577\n\
1578Create a new compressor object. This object may be used to compress\n\
1579data sequentially. If you want to compress data in one shot, use the\n\
1580compress() function instead. The compresslevel parameter, if given,\n\
1581must be a number between 1 and 9.\n\
1582");
1583
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001584static PyTypeObject BZ2Comp_Type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001585 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001586 "bz2.BZ2Compressor", /*tp_name*/
1587 sizeof(BZ2CompObject), /*tp_basicsize*/
1588 0, /*tp_itemsize*/
1589 (destructor)BZ2Comp_dealloc, /*tp_dealloc*/
1590 0, /*tp_print*/
1591 0, /*tp_getattr*/
1592 0, /*tp_setattr*/
1593 0, /*tp_compare*/
1594 0, /*tp_repr*/
1595 0, /*tp_as_number*/
1596 0, /*tp_as_sequence*/
1597 0, /*tp_as_mapping*/
1598 0, /*tp_hash*/
1599 0, /*tp_call*/
1600 0, /*tp_str*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001601 PyObject_GenericGetAttr,/*tp_getattro*/
1602 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001603 0, /*tp_as_buffer*/
1604 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1605 BZ2Comp__doc__, /*tp_doc*/
1606 0, /*tp_traverse*/
1607 0, /*tp_clear*/
1608 0, /*tp_richcompare*/
1609 0, /*tp_weaklistoffset*/
1610 0, /*tp_iter*/
1611 0, /*tp_iternext*/
1612 BZ2Comp_methods, /*tp_methods*/
1613 0, /*tp_members*/
1614 0, /*tp_getset*/
1615 0, /*tp_base*/
1616 0, /*tp_dict*/
1617 0, /*tp_descr_get*/
1618 0, /*tp_descr_set*/
1619 0, /*tp_dictoffset*/
1620 (initproc)BZ2Comp_init, /*tp_init*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001621 PyType_GenericAlloc, /*tp_alloc*/
1622 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001623 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001624 0, /*tp_is_gc*/
1625};
1626
1627
1628/* ===================================================================== */
1629/* Members of BZ2Decomp. */
1630
Gustavo Niemeyera33d0aa2003-02-11 18:46:20 +00001631#undef OFF
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001632#define OFF(x) offsetof(BZ2DecompObject, x)
1633
1634static PyMemberDef BZ2Decomp_members[] = {
Guido van Rossum33d26892007-08-05 15:29:28 +00001635 {"unused_data", T_OBJECT, OFF(unused_data), READONLY},
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001636 {NULL} /* Sentinel */
1637};
1638
1639
1640/* ===================================================================== */
1641/* Methods of BZ2Decomp. */
1642
1643PyDoc_STRVAR(BZ2Decomp_decompress__doc__,
1644"decompress(data) -> string\n\
1645\n\
1646Provide more data to the decompressor object. It will return chunks\n\
1647of decompressed data whenever possible. If you try to decompress data\n\
1648after the end of stream is found, EOFError will be raised. If any data\n\
1649was found after the end of stream, it'll be ignored and saved in\n\
1650unused_data attribute.\n\
1651");
1652
1653static PyObject *
1654BZ2Decomp_decompress(BZ2DecompObject *self, PyObject *args)
1655{
Martin v. Löwis423be952008-08-13 15:53:07 +00001656 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001657 char *data;
1658 int datasize;
1659 int bufsize = SMALLCHUNK;
Martin v. Löwisb9a0f912003-03-29 10:06:18 +00001660 PY_LONG_LONG totalout;
Neal Norwitz18142c02002-11-05 18:17:32 +00001661 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001662 bz_stream *bzs = &self->bzs;
1663 int bzerror;
1664
Martin v. Löwis423be952008-08-13 15:53:07 +00001665 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001666 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001667 data = pdata.buf;
1668 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001669
1670 ACQUIRE_LOCK(self);
1671 if (!self->running) {
1672 PyErr_SetString(PyExc_EOFError, "end of stream was "
1673 "already found");
1674 goto error;
1675 }
1676
Christian Heimes72b710a2008-05-26 13:28:38 +00001677 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001678 if (!ret)
1679 goto error;
1680
1681 bzs->next_in = data;
1682 bzs->avail_in = datasize;
1683 bzs->next_out = BUF(ret);
1684 bzs->avail_out = bufsize;
1685
1686 totalout = BZS_TOTAL_OUT(bzs);
1687
1688 for (;;) {
1689 Py_BEGIN_ALLOW_THREADS
1690 bzerror = BZ2_bzDecompress(bzs);
1691 Py_END_ALLOW_THREADS
1692 if (bzerror == BZ_STREAM_END) {
1693 if (bzs->avail_in != 0) {
1694 Py_DECREF(self->unused_data);
1695 self->unused_data =
Christian Heimes72b710a2008-05-26 13:28:38 +00001696 PyBytes_FromStringAndSize(bzs->next_in,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001697 bzs->avail_in);
1698 }
1699 self->running = 0;
1700 break;
1701 }
1702 if (bzerror != BZ_OK) {
1703 Util_CatchBZ2Error(bzerror);
1704 goto error;
1705 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00001706 if (bzs->avail_in == 0)
1707 break; /* no more input data */
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001708 if (bzs->avail_out == 0) {
1709 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001710 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001711 BZ2_bzDecompressEnd(bzs);
1712 goto error;
1713 }
1714 bzs->next_out = BUF(ret);
1715 bzs->next_out = BUF(ret) + (BZS_TOTAL_OUT(bzs)
1716 - totalout);
1717 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001718 }
1719 }
1720
Guido van Rossum522a6c62007-05-22 23:13:45 +00001721 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001722 if (_PyBytes_Resize(&ret,
Guido van Rossum522a6c62007-05-22 23:13:45 +00001723 (Py_ssize_t)(BZS_TOTAL_OUT(bzs) - totalout)) < 0)
1724 goto error;
1725 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001726
1727 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001728 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001729 return ret;
1730
1731error:
1732 RELEASE_LOCK(self);
Martin v. Löwis423be952008-08-13 15:53:07 +00001733 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001734 Py_XDECREF(ret);
1735 return NULL;
1736}
1737
1738static PyMethodDef BZ2Decomp_methods[] = {
1739 {"decompress", (PyCFunction)BZ2Decomp_decompress, METH_VARARGS, BZ2Decomp_decompress__doc__},
1740 {NULL, NULL} /* sentinel */
1741};
1742
1743
1744/* ===================================================================== */
1745/* Slot definitions for BZ2Decomp_Type. */
1746
1747static int
1748BZ2Decomp_init(BZ2DecompObject *self, PyObject *args, PyObject *kwargs)
1749{
1750 int bzerror;
1751
1752 if (!PyArg_ParseTuple(args, ":BZ2Decompressor"))
1753 return -1;
1754
1755#ifdef WITH_THREAD
1756 self->lock = PyThread_allocate_lock();
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001757 if (!self->lock) {
1758 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001759 goto error;
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001760 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001761#endif
1762
Christian Heimes72b710a2008-05-26 13:28:38 +00001763 self->unused_data = PyBytes_FromStringAndSize("", 0);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001764 if (!self->unused_data)
1765 goto error;
1766
1767 memset(&self->bzs, 0, sizeof(bz_stream));
1768 bzerror = BZ2_bzDecompressInit(&self->bzs, 0, 0);
1769 if (bzerror != BZ_OK) {
1770 Util_CatchBZ2Error(bzerror);
1771 goto error;
1772 }
1773
1774 self->running = 1;
1775
1776 return 0;
1777
1778error:
1779#ifdef WITH_THREAD
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001780 if (self->lock) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001781 PyThread_free_lock(self->lock);
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001782 self->lock = NULL;
1783 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001784#endif
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001785 Py_CLEAR(self->unused_data);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001786 return -1;
1787}
1788
1789static void
1790BZ2Decomp_dealloc(BZ2DecompObject *self)
1791{
1792#ifdef WITH_THREAD
1793 if (self->lock)
1794 PyThread_free_lock(self->lock);
1795#endif
1796 Py_XDECREF(self->unused_data);
1797 BZ2_bzDecompressEnd(&self->bzs);
Christian Heimes90aa7642007-12-19 02:45:37 +00001798 Py_TYPE(self)->tp_free((PyObject *)self);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001799}
1800
1801
1802/* ===================================================================== */
1803/* BZ2Decomp_Type definition. */
1804
1805PyDoc_STRVAR(BZ2Decomp__doc__,
1806"BZ2Decompressor() -> decompressor object\n\
1807\n\
1808Create a new decompressor object. This object may be used to decompress\n\
1809data sequentially. If you want to decompress data in one shot, use the\n\
1810decompress() function instead.\n\
1811");
1812
Gustavo Niemeyer49ea7be2002-11-08 14:31:49 +00001813static PyTypeObject BZ2Decomp_Type = {
Martin v. Löwis9f2e3462007-07-21 17:22:18 +00001814 PyVarObject_HEAD_INIT(NULL, 0)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001815 "bz2.BZ2Decompressor", /*tp_name*/
1816 sizeof(BZ2DecompObject), /*tp_basicsize*/
1817 0, /*tp_itemsize*/
1818 (destructor)BZ2Decomp_dealloc, /*tp_dealloc*/
1819 0, /*tp_print*/
1820 0, /*tp_getattr*/
1821 0, /*tp_setattr*/
1822 0, /*tp_compare*/
1823 0, /*tp_repr*/
1824 0, /*tp_as_number*/
1825 0, /*tp_as_sequence*/
1826 0, /*tp_as_mapping*/
1827 0, /*tp_hash*/
1828 0, /*tp_call*/
1829 0, /*tp_str*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001830 PyObject_GenericGetAttr,/*tp_getattro*/
1831 PyObject_GenericSetAttr,/*tp_setattro*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001832 0, /*tp_as_buffer*/
1833 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1834 BZ2Decomp__doc__, /*tp_doc*/
1835 0, /*tp_traverse*/
1836 0, /*tp_clear*/
1837 0, /*tp_richcompare*/
1838 0, /*tp_weaklistoffset*/
1839 0, /*tp_iter*/
1840 0, /*tp_iternext*/
1841 BZ2Decomp_methods, /*tp_methods*/
1842 BZ2Decomp_members, /*tp_members*/
1843 0, /*tp_getset*/
1844 0, /*tp_base*/
1845 0, /*tp_dict*/
1846 0, /*tp_descr_get*/
1847 0, /*tp_descr_set*/
1848 0, /*tp_dictoffset*/
1849 (initproc)BZ2Decomp_init, /*tp_init*/
Jason Tishlerfb8595d2003-01-06 12:41:26 +00001850 PyType_GenericAlloc, /*tp_alloc*/
1851 PyType_GenericNew, /*tp_new*/
Neal Norwitz30d1c512007-08-19 22:48:23 +00001852 PyObject_Free, /*tp_free*/
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001853 0, /*tp_is_gc*/
1854};
1855
1856
1857/* ===================================================================== */
1858/* Module functions. */
1859
1860PyDoc_STRVAR(bz2_compress__doc__,
1861"compress(data [, compresslevel=9]) -> string\n\
1862\n\
1863Compress data in one shot. If you want to compress data sequentially,\n\
1864use an instance of BZ2Compressor instead. The compresslevel parameter, if\n\
1865given, must be a number between 1 and 9.\n\
1866");
1867
1868static PyObject *
1869bz2_compress(PyObject *self, PyObject *args, PyObject *kwargs)
1870{
1871 int compresslevel=9;
Martin v. Löwis423be952008-08-13 15:53:07 +00001872 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001873 char *data;
1874 int datasize;
1875 int bufsize;
Gustavo Niemeyer7d7930b2002-11-05 18:41:53 +00001876 PyObject *ret = NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001877 bz_stream _bzs;
1878 bz_stream *bzs = &_bzs;
1879 int bzerror;
Martin v. Löwis15e62742006-02-27 16:46:16 +00001880 static char *kwlist[] = {"data", "compresslevel", 0};
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001881
Martin v. Löwis423be952008-08-13 15:53:07 +00001882 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "y*|i",
1883 kwlist, &pdata,
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001884 &compresslevel))
1885 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001886 data = pdata.buf;
1887 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001888
1889 if (compresslevel < 1 || compresslevel > 9) {
1890 PyErr_SetString(PyExc_ValueError,
1891 "compresslevel must be between 1 and 9");
Martin v. Löwis423be952008-08-13 15:53:07 +00001892 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001893 return NULL;
1894 }
1895
1896 /* Conforming to bz2 manual, this is large enough to fit compressed
1897 * data in one shot. We will check it later anyway. */
1898 bufsize = datasize + (datasize/100+1) + 600;
1899
Christian Heimes72b710a2008-05-26 13:28:38 +00001900 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Martin v. Löwis423be952008-08-13 15:53:07 +00001901 if (!ret) {
1902 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001903 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001904 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001905
1906 memset(bzs, 0, sizeof(bz_stream));
1907
1908 bzs->next_in = data;
1909 bzs->avail_in = datasize;
1910 bzs->next_out = BUF(ret);
1911 bzs->avail_out = bufsize;
1912
1913 bzerror = BZ2_bzCompressInit(bzs, compresslevel, 0, 0);
1914 if (bzerror != BZ_OK) {
1915 Util_CatchBZ2Error(bzerror);
Martin v. Löwis423be952008-08-13 15:53:07 +00001916 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001917 Py_DECREF(ret);
1918 return NULL;
1919 }
Tim Peterse3228092002-11-09 04:21:44 +00001920
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001921 for (;;) {
1922 Py_BEGIN_ALLOW_THREADS
1923 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1924 Py_END_ALLOW_THREADS
1925 if (bzerror == BZ_STREAM_END) {
1926 break;
1927 } else if (bzerror != BZ_FINISH_OK) {
1928 BZ2_bzCompressEnd(bzs);
1929 Util_CatchBZ2Error(bzerror);
Martin v. Löwis423be952008-08-13 15:53:07 +00001930 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001931 Py_DECREF(ret);
1932 return NULL;
1933 }
1934 if (bzs->avail_out == 0) {
1935 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00001936 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001937 BZ2_bzCompressEnd(bzs);
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 bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
1942 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
1943 }
1944 }
1945
Guido van Rossum522a6c62007-05-22 23:13:45 +00001946 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00001947 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +00001948 ret = NULL;
1949 }
1950 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001951 BZ2_bzCompressEnd(bzs);
1952
Martin v. Löwis423be952008-08-13 15:53:07 +00001953 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001954 return ret;
1955}
1956
1957PyDoc_STRVAR(bz2_decompress__doc__,
1958"decompress(data) -> decompressed data\n\
1959\n\
1960Decompress data in one shot. If you want to decompress data sequentially,\n\
1961use an instance of BZ2Decompressor instead.\n\
1962");
1963
1964static PyObject *
1965bz2_decompress(PyObject *self, PyObject *args)
1966{
Martin v. Löwis423be952008-08-13 15:53:07 +00001967 Py_buffer pdata;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001968 char *data;
1969 int datasize;
1970 int bufsize = SMALLCHUNK;
1971 PyObject *ret;
1972 bz_stream _bzs;
1973 bz_stream *bzs = &_bzs;
1974 int bzerror;
1975
Martin v. Löwis423be952008-08-13 15:53:07 +00001976 if (!PyArg_ParseTuple(args, "y*:decompress", &pdata))
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001977 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001978 data = pdata.buf;
1979 datasize = pdata.len;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001980
Martin v. Löwis423be952008-08-13 15:53:07 +00001981 if (datasize == 0) {
1982 PyBuffer_Release(&pdata);
Christian Heimes72b710a2008-05-26 13:28:38 +00001983 return PyBytes_FromStringAndSize("", 0);
Martin v. Löwis423be952008-08-13 15:53:07 +00001984 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001985
Christian Heimes72b710a2008-05-26 13:28:38 +00001986 ret = PyBytes_FromStringAndSize(NULL, bufsize);
Martin v. Löwis423be952008-08-13 15:53:07 +00001987 if (!ret) {
1988 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001989 return NULL;
Martin v. Löwis423be952008-08-13 15:53:07 +00001990 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00001991
1992 memset(bzs, 0, sizeof(bz_stream));
1993
1994 bzs->next_in = data;
1995 bzs->avail_in = datasize;
1996 bzs->next_out = BUF(ret);
1997 bzs->avail_out = bufsize;
1998
1999 bzerror = BZ2_bzDecompressInit(bzs, 0, 0);
2000 if (bzerror != BZ_OK) {
2001 Util_CatchBZ2Error(bzerror);
2002 Py_DECREF(ret);
Martin v. Löwis423be952008-08-13 15:53:07 +00002003 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002004 return NULL;
2005 }
Tim Peterse3228092002-11-09 04:21:44 +00002006
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002007 for (;;) {
2008 Py_BEGIN_ALLOW_THREADS
2009 bzerror = BZ2_bzDecompress(bzs);
2010 Py_END_ALLOW_THREADS
2011 if (bzerror == BZ_STREAM_END) {
2012 break;
2013 } else if (bzerror != BZ_OK) {
2014 BZ2_bzDecompressEnd(bzs);
2015 Util_CatchBZ2Error(bzerror);
Martin v. Löwis423be952008-08-13 15:53:07 +00002016 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002017 Py_DECREF(ret);
2018 return NULL;
2019 }
Guido van Rossumd8faa362007-04-27 19:54:29 +00002020 if (bzs->avail_in == 0) {
2021 BZ2_bzDecompressEnd(bzs);
2022 PyErr_SetString(PyExc_ValueError,
2023 "couldn't find end of stream");
Martin v. Löwis423be952008-08-13 15:53:07 +00002024 PyBuffer_Release(&pdata);
Guido van Rossumd8faa362007-04-27 19:54:29 +00002025 Py_DECREF(ret);
2026 return NULL;
2027 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002028 if (bzs->avail_out == 0) {
2029 bufsize = Util_NewBufferSize(bufsize);
Christian Heimes72b710a2008-05-26 13:28:38 +00002030 if (_PyBytes_Resize(&ret, bufsize) < 0) {
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002031 BZ2_bzDecompressEnd(bzs);
Martin v. Löwis423be952008-08-13 15:53:07 +00002032 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002033 return NULL;
2034 }
2035 bzs->next_out = BUF(ret) + BZS_TOTAL_OUT(bzs);
2036 bzs->avail_out = bufsize - (bzs->next_out - BUF(ret));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002037 }
2038 }
2039
Guido van Rossum522a6c62007-05-22 23:13:45 +00002040 if (bzs->avail_out != 0) {
Christian Heimes72b710a2008-05-26 13:28:38 +00002041 if (_PyBytes_Resize(&ret, (Py_ssize_t)BZS_TOTAL_OUT(bzs)) < 0) {
Guido van Rossum522a6c62007-05-22 23:13:45 +00002042 ret = NULL;
2043 }
2044 }
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002045 BZ2_bzDecompressEnd(bzs);
Martin v. Löwis423be952008-08-13 15:53:07 +00002046 PyBuffer_Release(&pdata);
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002047
2048 return ret;
2049}
2050
2051static PyMethodDef bz2_methods[] = {
2052 {"compress", (PyCFunction) bz2_compress, METH_VARARGS|METH_KEYWORDS,
2053 bz2_compress__doc__},
2054 {"decompress", (PyCFunction) bz2_decompress, METH_VARARGS,
2055 bz2_decompress__doc__},
2056 {NULL, NULL} /* sentinel */
2057};
2058
2059/* ===================================================================== */
2060/* Initialization function. */
2061
2062PyDoc_STRVAR(bz2__doc__,
2063"The python bz2 module provides a comprehensive interface for\n\
2064the bz2 compression library. It implements a complete file\n\
2065interface, one shot (de)compression functions, and types for\n\
2066sequential (de)compression.\n\
2067");
2068
Martin v. Löwis1a214512008-06-11 05:26:20 +00002069
2070static struct PyModuleDef bz2module = {
2071 PyModuleDef_HEAD_INIT,
2072 "bz2",
2073 bz2__doc__,
2074 -1,
2075 bz2_methods,
2076 NULL,
2077 NULL,
2078 NULL,
2079 NULL
2080};
2081
Neal Norwitz21d896c2003-07-01 20:15:21 +00002082PyMODINIT_FUNC
Martin v. Löwis1a214512008-06-11 05:26:20 +00002083PyInit_bz2(void)
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002084{
2085 PyObject *m;
2086
Christian Heimes90aa7642007-12-19 02:45:37 +00002087 Py_TYPE(&BZ2File_Type) = &PyType_Type;
2088 Py_TYPE(&BZ2Comp_Type) = &PyType_Type;
2089 Py_TYPE(&BZ2Decomp_Type) = &PyType_Type;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002090
Martin v. Löwis1a214512008-06-11 05:26:20 +00002091 m = PyModule_Create(&bz2module);
Neal Norwitz1ac754f2006-01-19 06:09:39 +00002092 if (m == NULL)
Martin v. Löwis1a214512008-06-11 05:26:20 +00002093 return NULL;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002094
Neal Norwitz53cbdaa2007-08-23 21:42:55 +00002095 PyModule_AddObject(m, "__author__", PyUnicode_FromString(__author__));
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002096
2097 Py_INCREF(&BZ2File_Type);
2098 PyModule_AddObject(m, "BZ2File", (PyObject *)&BZ2File_Type);
2099
2100 Py_INCREF(&BZ2Comp_Type);
2101 PyModule_AddObject(m, "BZ2Compressor", (PyObject *)&BZ2Comp_Type);
2102
2103 Py_INCREF(&BZ2Decomp_Type);
2104 PyModule_AddObject(m, "BZ2Decompressor", (PyObject *)&BZ2Decomp_Type);
Martin v. Löwis1a214512008-06-11 05:26:20 +00002105 return m;
Gustavo Niemeyerf8ca8362002-11-05 16:50:05 +00002106}