1 /*
2
3 python-bz2 - python bz2 library interface
4
5 Copyright (c) 2002 Gustavo Niemeyer <niemeyer@conectiva.com>
6 Copyright (c) 2002 Python Software Foundation; All Rights Reserved
7
8 */
9
10 #include "Python.h"
11 #include <stdio.h>
12 #include <bzlib.h>
13 #include "structmember.h"
14
15 #ifdef WITH_THREAD
16 #include "pythread.h"
17 #endif
18
19 static char __author__[] =
20 "The bz2 python module was written by:\n\
21 \n\
22 Gustavo Niemeyer <niemeyer@conectiva.com>\n\
23 ";
24
25 /* Our very own off_t-like type, 64-bit if possible */
26 /* copied from Objects/fileobject.c */
27 #if !defined(HAVE_LARGEFILE_SUPPORT)
28 typedef off_t Py_off_t;
29 #elif SIZEOF_OFF_T >= 8
30 typedef off_t Py_off_t;
31 #elif SIZEOF_FPOS_T >= 8
32 typedef 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
37 #define BUF(v) PyString_AS_STRING((PyStringObject *)v)
38
39 #define MODE_CLOSED 0
40 #define MODE_READ 1
41 #define MODE_READ_EOF 2
42 #define MODE_WRITE 3
43
44
45 #ifndef BZ_CONFIG_ERROR
46
47 #define BZ2_bzRead bzRead
48 #define BZ2_bzReadOpen bzReadOpen
49 #define BZ2_bzReadClose bzReadClose
50 #define BZ2_bzWrite bzWrite
51 #define BZ2_bzWriteOpen bzWriteOpen
52 #define BZ2_bzWriteClose bzWriteClose
53 #define BZ2_bzCompress bzCompress
54 #define BZ2_bzCompressInit bzCompressInit
55 #define BZ2_bzCompressEnd bzCompressEnd
56 #define BZ2_bzDecompress bzDecompress
57 #define BZ2_bzDecompressInit bzDecompressInit
58 #define BZ2_bzDecompressEnd bzDecompressEnd
59
60 #endif /* ! BZ_CONFIG_ERROR */
61
62
63 #ifdef WITH_THREAD
64 #define ACQUIRE_LOCK(obj) do { \
65 if (!PyThread_acquire_lock(obj->lock, 0)) { \
66 Py_BEGIN_ALLOW_THREADS \
67 PyThread_acquire_lock(obj->lock, 1); \
68 Py_END_ALLOW_THREADS \
69 } } while(0)
70 #define RELEASE_LOCK(obj) PyThread_release_lock(obj->lock)
71 #else
72 #define ACQUIRE_LOCK(obj)
73 #define RELEASE_LOCK(obj)
74 #endif
75
76 #define MIN(X, Y) (((X) < (Y)) ? (X) : (Y))
77
78 /* Bits in f_newlinetypes */
79 #define NEWLINE_UNKNOWN 0 /* No newline seen, yet */
80 #define NEWLINE_CR 1 /* \r newline seen */
81 #define NEWLINE_LF 2 /* \n newline seen */
82 #define NEWLINE_CRLF 4 /* \r\n newline seen */
83
84 /* ===================================================================== */
85 /* Structure definitions. */
86
87 typedef struct {
88 PyObject_HEAD
89 PyObject *file;
90
91 char* f_buf; /* Allocated readahead buffer */
92 char* f_bufend; /* Points after last occupied position */
93 char* f_bufptr; /* Current buffer position */
94
95 int f_softspace; /* Flag used by 'print' command */
96
97 int f_univ_newline; /* Handle any newline convention */
98 int f_newlinetypes; /* Types of newlines seen */
99 int f_skipnextlf; /* Skip next \n */
100
101 BZFILE *fp;
102 int mode;
103 Py_off_t pos;
104 Py_off_t size;
105 #ifdef WITH_THREAD
106 PyThread_type_lock lock;
107 #endif
108 } BZ2FileObject;
109
110 typedef struct {
111 PyObject_HEAD
112 bz_stream bzs;
113 int running;
114 #ifdef WITH_THREAD
115 PyThread_type_lock lock;
116 #endif
117 } BZ2CompObject;
118
119 typedef struct {
120 PyObject_HEAD
121 bz_stream bzs;
122 int running;
123 PyObject *unused_data;
124 #ifdef WITH_THREAD
125 PyThread_type_lock lock;
126 #endif
127 } BZ2DecompObject;
128
129 /* ===================================================================== */
130 /* Utility functions. */
131
132 /* Refuse regular I/O if there's data in the iteration-buffer.
133 * Mixing them would cause data to arrive out of order, as the read*
134 * methods don't use the iteration buffer. */
135 static int
check_iterbuffered(BZ2FileObject * f)136 check_iterbuffered(BZ2FileObject *f)
137 {
138 if (f->f_buf != NULL &&
139 (f->f_bufend - f->f_bufptr) > 0 &&
140 f->f_buf[0] != '\0') {
141 PyErr_SetString(PyExc_ValueError,
142 "Mixing iteration and read methods would lose data");
143 return -1;
144 }
145 return 0;
146 }
147
148 static int
Util_CatchBZ2Error(int bzerror)149 Util_CatchBZ2Error(int bzerror)
150 {
151 int ret = 0;
152 switch(bzerror) {
153 case BZ_OK:
154 case BZ_STREAM_END:
155 break;
156
157 #ifdef BZ_CONFIG_ERROR
158 case BZ_CONFIG_ERROR:
159 PyErr_SetString(PyExc_SystemError,
160 "the bz2 library was not compiled "
161 "correctly");
162 ret = 1;
163 break;
164 #endif
165
166 case BZ_PARAM_ERROR:
167 PyErr_SetString(PyExc_ValueError,
168 "the bz2 library has received wrong "
169 "parameters");
170 ret = 1;
171 break;
172
173 case BZ_MEM_ERROR:
174 PyErr_NoMemory();
175 ret = 1;
176 break;
177
178 case BZ_DATA_ERROR:
179 case BZ_DATA_ERROR_MAGIC:
180 PyErr_SetString(PyExc_IOError, "invalid data stream");
181 ret = 1;
182 break;
183
184 case BZ_IO_ERROR:
185 PyErr_SetString(PyExc_IOError, "unknown IO error");
186 ret = 1;
187 break;
188
189 case BZ_UNEXPECTED_EOF:
190 PyErr_SetString(PyExc_EOFError,
191 "compressed file ended before the "
192 "logical end-of-stream was detected");
193 ret = 1;
194 break;
195
196 case BZ_SEQUENCE_ERROR:
197 PyErr_SetString(PyExc_RuntimeError,
198 "wrong sequence of bz2 library "
199 "commands used");
200 ret = 1;
201 break;
202 }
203 return ret;
204 }
205
206 #if BUFSIZ < 8192
207 #define SMALLCHUNK 8192
208 #else
209 #define SMALLCHUNK BUFSIZ
210 #endif
211
212 /* This is a hacked version of Python's fileobject.c:new_buffersize(). */
213 static size_t
Util_NewBufferSize(size_t currentsize)214 Util_NewBufferSize(size_t currentsize)
215 {
216 /* Expand the buffer by an amount proportional to the current size,
217 giving us amortized linear-time behavior. Use a less-than-double
218 growth factor to avoid excessive allocation. */
219 return currentsize + (currentsize >> 3) + 6;
220 }
221
222 static int
Util_GrowBuffer(PyObject ** buf)223 Util_GrowBuffer(PyObject **buf)
224 {
225 size_t size = PyString_GET_SIZE(*buf);
226 size_t new_size = Util_NewBufferSize(size);
227 if (new_size > size) {
228 return _PyString_Resize(buf, new_size);
229 } else { /* overflow */
230 PyErr_SetString(PyExc_OverflowError,
231 "Unable to allocate buffer - output too large");
232 return -1;
233 }
234 }
235
236 /* This is a hacked version of Python's fileobject.c:get_line(). */
237 static PyObject *
Util_GetLine(BZ2FileObject * f,int n)238 Util_GetLine(BZ2FileObject *f, int n)
239 {
240 char c;
241 char *buf, *end;
242 size_t total_v_size; /* total # of slots in buffer */
243 size_t used_v_size; /* # used slots in buffer */
244 size_t increment; /* amount to increment the buffer */
245 PyObject *v;
246 int bzerror;
247 int bytes_read;
248 int newlinetypes = f->f_newlinetypes;
249 int skipnextlf = f->f_skipnextlf;
250 int univ_newline = f->f_univ_newline;
251
252 total_v_size = n > 0 ? n : 100;
253 v = PyString_FromStringAndSize((char *)NULL, total_v_size);
254 if (v == NULL)
255 return NULL;
256
257 buf = BUF(v);
258 end = buf + total_v_size;
259
260 for (;;) {
261 Py_BEGIN_ALLOW_THREADS
262 while (buf != end) {
263 bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
264 f->pos++;
265 if (bytes_read == 0) break;
266 if (univ_newline) {
267 if (skipnextlf) {
268 skipnextlf = 0;
269 if (c == '\n') {
270 /* Seeing a \n here with skipnextlf true means we
271 * saw a \r before.
272 */
273 newlinetypes |= NEWLINE_CRLF;
274 if (bzerror != BZ_OK) break;
275 bytes_read = BZ2_bzRead(&bzerror, f->fp, &c, 1);
276 f->pos++;
277 if (bytes_read == 0) break;
278 } else {
279 newlinetypes |= NEWLINE_CR;
280 }
281 }
282 if (c == '\r') {
283 skipnextlf = 1;
284 c = '\n';
285 } else if (c == '\n')
286 newlinetypes |= NEWLINE_LF;
287 }
288 *buf++ = c;
289 if (bzerror != BZ_OK || c == '\n') break;
290 }
291 if (univ_newline && bzerror == BZ_STREAM_END && skipnextlf)
292 newlinetypes |= NEWLINE_CR;
293 Py_END_ALLOW_THREADS
294 f->f_newlinetypes = newlinetypes;
295 f->f_skipnextlf = skipnextlf;
296 if (bzerror == BZ_STREAM_END) {
297 f->size = f->pos;
298 f->mode = MODE_READ_EOF;
299 break;
300 } else if (bzerror != BZ_OK) {
301 Util_CatchBZ2Error(bzerror);
302 Py_DECREF(v);
303 return NULL;
304 }
305 if (c == '\n')
306 break;
307 /* Must be because buf == end */
308 if (n > 0)
309 break;
310 used_v_size = total_v_size;
311 increment = total_v_size >> 2; /* mild exponential growth */
312 total_v_size += increment;
313 if (total_v_size > INT_MAX) {
314 PyErr_SetString(PyExc_OverflowError,
315 "line is longer than a Python string can hold");
316 Py_DECREF(v);
317 return NULL;
318 }
319 if (_PyString_Resize(&v, total_v_size) < 0)
320 return NULL;
321 buf = BUF(v) + used_v_size;
322 end = BUF(v) + total_v_size;
323 }
324
325 used_v_size = buf - BUF(v);
326 if (used_v_size != total_v_size)
327 _PyString_Resize(&v, used_v_size);
328 return v;
329 }
330
331 /* This is a hacked version of Python's
332 * fileobject.c:Py_UniversalNewlineFread(). */
333 size_t
Util_UnivNewlineRead(int * bzerror,BZFILE * stream,char * buf,size_t n,BZ2FileObject * f)334 Util_UnivNewlineRead(int *bzerror, BZFILE *stream,
335 char* buf, size_t n, BZ2FileObject *f)
336 {
337 char *dst = buf;
338 int newlinetypes, skipnextlf;
339
340 assert(buf != NULL);
341 assert(stream != NULL);
342
343 if (!f->f_univ_newline)
344 return BZ2_bzRead(bzerror, stream, buf, n);
345
346 newlinetypes = f->f_newlinetypes;
347 skipnextlf = f->f_skipnextlf;
348
349 /* Invariant: n is the number of bytes remaining to be filled
350 * in the buffer.
351 */
352 while (n) {
353 size_t nread;
354 int shortread;
355 char *src = dst;
356
357 nread = BZ2_bzRead(bzerror, stream, dst, n);
358 assert(nread <= n);
359 n -= nread; /* assuming 1 byte out for each in; will adjust */
360 shortread = n != 0; /* true iff EOF or error */
361 while (nread--) {
362 char c = *src++;
363 if (c == '\r') {
364 /* Save as LF and set flag to skip next LF. */
365 *dst++ = '\n';
366 skipnextlf = 1;
367 }
368 else if (skipnextlf && c == '\n') {
369 /* Skip LF, and remember we saw CR LF. */
370 skipnextlf = 0;
371 newlinetypes |= NEWLINE_CRLF;
372 ++n;
373 }
374 else {
375 /* Normal char to be stored in buffer. Also
376 * update the newlinetypes flag if either this
377 * is an LF or the previous char was a CR.
378 */
379 if (c == '\n')
380 newlinetypes |= NEWLINE_LF;
381 else if (skipnextlf)
382 newlinetypes |= NEWLINE_CR;
383 *dst++ = c;
384 skipnextlf = 0;
385 }
386 }
387 if (shortread) {
388 /* If this is EOF, update type flags. */
389 if (skipnextlf && *bzerror == BZ_STREAM_END)
390 newlinetypes |= NEWLINE_CR;
391 break;
392 }
393 }
394 f->f_newlinetypes = newlinetypes;
395 f->f_skipnextlf = skipnextlf;
396 return dst - buf;
397 }
398
399 /* This is a hacked version of Python's fileobject.c:drop_readahead(). */
400 static void
Util_DropReadAhead(BZ2FileObject * f)401 Util_DropReadAhead(BZ2FileObject *f)
402 {
403 if (f->f_buf != NULL) {
404 PyMem_Free(f->f_buf);
405 f->f_buf = NULL;
406 }
407 }
408
409 /* This is a hacked version of Python's fileobject.c:readahead(). */
410 static int
Util_ReadAhead(BZ2FileObject * f,int bufsize)411 Util_ReadAhead(BZ2FileObject *f, int bufsize)
412 {
413 int chunksize;
414 int bzerror;
415
416 if (f->f_buf != NULL) {
417 if((f->f_bufend - f->f_bufptr) >= 1)
418 return 0;
419 else
420 Util_DropReadAhead(f);
421 }
422 if (f->mode == MODE_READ_EOF) {
423 f->f_bufptr = f->f_buf;
424 f->f_bufend = f->f_buf;
425 return 0;
426 }
427 if ((f->f_buf = PyMem_Malloc(bufsize)) == NULL) {
428 PyErr_NoMemory();
429 return -1;
430 }
431 Py_BEGIN_ALLOW_THREADS
432 chunksize = Util_UnivNewlineRead(&bzerror, f->fp, f->f_buf,
433 bufsize, f);
434 Py_END_ALLOW_THREADS
435 f->pos += chunksize;
436 if (bzerror == BZ_STREAM_END) {
437 f->size = f->pos;
438 f->mode = MODE_READ_EOF;
439 } else if (bzerror != BZ_OK) {
440 Util_CatchBZ2Error(bzerror);
441 Util_DropReadAhead(f);
442 return -1;
443 }
444 f->f_bufptr = f->f_buf;
445 f->f_bufend = f->f_buf + chunksize;
446 return 0;
447 }
448
449 /* This is a hacked version of Python's
450 * fileobject.c:readahead_get_line_skip(). */
451 static PyStringObject *
Util_ReadAheadGetLineSkip(BZ2FileObject * f,int skip,int bufsize)452 Util_ReadAheadGetLineSkip(BZ2FileObject *f, int skip, int bufsize)
453 {
454 PyStringObject* s;
455 char *bufptr;
456 char *buf;
457 int len;
458
459 if (f->f_buf == NULL)
460 if (Util_ReadAhead(f, bufsize) < 0)
461 return NULL;
462
463 len = f->f_bufend - f->f_bufptr;
464 if (len == 0)
465 return (PyStringObject *)
466 PyString_FromStringAndSize(NULL, skip);
467 bufptr = memchr(f->f_bufptr, '\n', len);
468 if (bufptr != NULL) {
469 bufptr++; /* Count the '\n' */
470 len = bufptr - f->f_bufptr;
471 s = (PyStringObject *)
472 PyString_FromStringAndSize(NULL, skip+len);
473 if (s == NULL)
474 return NULL;
475 memcpy(PyString_AS_STRING(s)+skip, f->f_bufptr, len);
476 f->f_bufptr = bufptr;
477 if (bufptr == f->f_bufend)
478 Util_DropReadAhead(f);
479 } else {
480 bufptr = f->f_bufptr;
481 buf = f->f_buf;
482 f->f_buf = NULL; /* Force new readahead buffer */
483 s = Util_ReadAheadGetLineSkip(f, skip+len,
484 bufsize + (bufsize>>2));
485 if (s == NULL) {
486 PyMem_Free(buf);
487 return NULL;
488 }
489 memcpy(PyString_AS_STRING(s)+skip, bufptr, len);
490 PyMem_Free(buf);
491 }
492 return s;
493 }
494
495 /* ===================================================================== */
496 /* Methods of BZ2File. */
497
498 PyDoc_STRVAR(BZ2File_read__doc__,
499 "read([size]) -> string\n\
500 \n\
501 Read at most size uncompressed bytes, returned as a string. If the size\n\
502 argument is negative or omitted, read until EOF is reached.\n\
503 ");
504
505 /* This is a hacked version of Python's fileobject.c:file_read(). */
506 static PyObject *
BZ2File_read(BZ2FileObject * self,PyObject * args)507 BZ2File_read(BZ2FileObject *self, PyObject *args)
508 {
509 long bytesrequested = -1;
510 size_t bytesread, buffersize, chunksize;
511 int bzerror;
512 PyObject *ret = NULL;
513
514 if (!PyArg_ParseTuple(args, "|l:read", &bytesrequested))
515 return NULL;
516
517 ACQUIRE_LOCK(self);
518 switch (self->mode) {
519 case MODE_READ:
520 break;
521 case MODE_READ_EOF:
522 ret = PyString_FromString("");
523 goto cleanup;
524 case MODE_CLOSED:
525 PyErr_SetString(PyExc_ValueError,
526 "I/O operation on closed file");
527 goto cleanup;
528 default:
529 PyErr_SetString(PyExc_IOError,
530 "file is not ready for reading");
531 goto cleanup;
532 }
533
534 /* refuse to mix with f.next() */
535 if (check_iterbuffered(self))
536 goto cleanup;
537
538 if (bytesrequested < 0)
539 buffersize = Util_NewBufferSize((size_t)0);
540 else
541 buffersize = bytesrequested;
542 if (buffersize > INT_MAX) {
543 PyErr_SetString(PyExc_OverflowError,
544 "requested number of bytes is "
545 "more than a Python string can hold");
546 goto cleanup;
547 }
548 ret = PyString_FromStringAndSize((char *)NULL, buffersize);
549 if (ret == NULL)
550 goto cleanup;
551 bytesread = 0;
552
553 for (;;) {
554 Py_BEGIN_ALLOW_THREADS
555 chunksize = Util_UnivNewlineRead(&bzerror, self->fp,
556 BUF(ret)+bytesread,
557 buffersize-bytesread,
558 self);
559 self->pos += chunksize;
560 Py_END_ALLOW_THREADS
561 bytesread += chunksize;
562 if (bzerror == BZ_STREAM_END) {
563 self->size = self->pos;
564 self->mode = MODE_READ_EOF;
565 break;
566 } else if (bzerror != BZ_OK) {
567 Util_CatchBZ2Error(bzerror);
568 Py_DECREF(ret);
569 ret = NULL;
570 goto cleanup;
571 }
572 if (bytesrequested < 0) {
573 buffersize = Util_NewBufferSize(buffersize);
574 if (_PyString_Resize(&ret, buffersize) < 0)
575 goto cleanup;
576 } else {
577 break;
578 }
579 }
580 if (bytesread != buffersize)
581 _PyString_Resize(&ret, bytesread);
582
583 cleanup:
584 RELEASE_LOCK(self);
585 return ret;
586 }
587
588 PyDoc_STRVAR(BZ2File_readline__doc__,
589 "readline([size]) -> string\n\
590 \n\
591 Return the next line from the file, as a string, retaining newline.\n\
592 A non-negative size argument will limit the maximum number of bytes to\n\
593 return (an incomplete line may be returned then). Return an empty\n\
594 string at EOF.\n\
595 ");
596
597 static PyObject *
BZ2File_readline(BZ2FileObject * self,PyObject * args)598 BZ2File_readline(BZ2FileObject *self, PyObject *args)
599 {
600 PyObject *ret = NULL;
601 int sizehint = -1;
602
603 if (!PyArg_ParseTuple(args, "|i:readline", &sizehint))
604 return NULL;
605
606 ACQUIRE_LOCK(self);
607 switch (self->mode) {
608 case MODE_READ:
609 break;
610 case MODE_READ_EOF:
611 ret = PyString_FromString("");
612 goto cleanup;
613 case MODE_CLOSED:
614 PyErr_SetString(PyExc_ValueError,
615 "I/O operation on closed file");
616 goto cleanup;
617 default:
618 PyErr_SetString(PyExc_IOError,
619 "file is not ready for reading");
620 goto cleanup;
621 }
622
623 /* refuse to mix with f.next() */
624 if (check_iterbuffered(self))
625 goto cleanup;
626
627 if (sizehint == 0)
628 ret = PyString_FromString("");
629 else
630 ret = Util_GetLine(self, (sizehint < 0) ? 0 : sizehint);
631
632 cleanup:
633 RELEASE_LOCK(self);
634 return ret;
635 }
636
637 PyDoc_STRVAR(BZ2File_readlines__doc__,
638 "readlines([size]) -> list\n\
639 \n\
640 Call readline() repeatedly and return a list of lines read.\n\
641 The optional size argument, if given, is an approximate bound on the\n\
642 total number of bytes in the lines returned.\n\
643 ");
644
645 /* This is a hacked version of Python's fileobject.c:file_readlines(). */
646 static PyObject *
BZ2File_readlines(BZ2FileObject * self,PyObject * args)647 BZ2File_readlines(BZ2FileObject *self, PyObject *args)
648 {
649 long sizehint = 0;
650 PyObject *list = NULL;
651 PyObject *line;
652 char small_buffer[SMALLCHUNK];
653 char *buffer = small_buffer;
654 size_t buffersize = SMALLCHUNK;
655 PyObject *big_buffer = NULL;
656 size_t nfilled = 0;
657 size_t nread;
658 size_t totalread = 0;
659 char *p, *q, *end;
660 int err;
661 int shortread = 0;
662 int bzerror;
663
664 if (!PyArg_ParseTuple(args, "|l:readlines", &sizehint))
665 return NULL;
666
667 ACQUIRE_LOCK(self);
668 switch (self->mode) {
669 case MODE_READ:
670 break;
671 case MODE_READ_EOF:
672 list = PyList_New(0);
673 goto cleanup;
674 case MODE_CLOSED:
675 PyErr_SetString(PyExc_ValueError,
676 "I/O operation on closed file");
677 goto cleanup;
678 default:
679 PyErr_SetString(PyExc_IOError,
680 "file is not ready for reading");
681 goto cleanup;
682 }
683
684 /* refuse to mix with f.next() */
685 if (check_iterbuffered(self))
686 goto cleanup;
687
688 if ((list = PyList_New(0)) == NULL)
689 goto cleanup;
690
691 for (;;) {
692 Py_BEGIN_ALLOW_THREADS
693 nread = Util_UnivNewlineRead(&bzerror, self->fp,
694 buffer+nfilled,
695 buffersize-nfilled, self);
696 self->pos += nread;
697 Py_END_ALLOW_THREADS
698 if (bzerror == BZ_STREAM_END) {
699 self->size = self->pos;
700 self->mode = MODE_READ_EOF;
701 if (nread == 0) {
702 sizehint = 0;
703 break;
704 }
705 shortread = 1;
706 } else if (bzerror != BZ_OK) {
707 Util_CatchBZ2Error(bzerror);
708 error:
709 Py_DECREF(list);
710 list = NULL;
711 goto cleanup;
712 }
713 totalread += nread;
714 p = memchr(buffer+nfilled, '\n', nread);
715 if (!shortread && p == NULL) {
716 /* Need a larger buffer to fit this line */
717 nfilled += nread;
718 buffersize *= 2;
719 if (buffersize > INT_MAX) {
720 PyErr_SetString(PyExc_OverflowError,
721 "line is longer than a Python string can hold");
722 goto error;
723 }
724 if (big_buffer == NULL) {
725 /* Create the big buffer */
726 big_buffer = PyString_FromStringAndSize(
727 NULL, buffersize);
728 if (big_buffer == NULL)
729 goto error;
730 buffer = PyString_AS_STRING(big_buffer);
731 memcpy(buffer, small_buffer, nfilled);
732 }
733 else {
734 /* Grow the big buffer */
735 if (_PyString_Resize(&big_buffer, buffersize))
736 goto error;
737 buffer = PyString_AS_STRING(big_buffer);
738 }
739 continue;
740 }
741 end = buffer+nfilled+nread;
742 q = buffer;
743 while (p != NULL) {
744 /* Process complete lines */
745 p++;
746 line = PyString_FromStringAndSize(q, p-q);
747 if (line == NULL)
748 goto error;
749 err = PyList_Append(list, line);
750 Py_DECREF(line);
751 if (err != 0)
752 goto error;
753 q = p;
754 p = memchr(q, '\n', end-q);
755 }
756 /* Move the remaining incomplete line to the start */
757 nfilled = end-q;
758 memmove(buffer, q, nfilled);
759 if (sizehint > 0)
760 if (totalread >= (size_t)sizehint)
761 break;
762 if (shortread) {
763 sizehint = 0;
764 break;
765 }
766 }
767 if (nfilled != 0) {
768 /* Partial last line */
769 line = PyString_FromStringAndSize(buffer, nfilled);
770 if (line == NULL)
771 goto error;
772 if (sizehint > 0) {
773 /* Need to complete the last line */
774 PyObject *rest = Util_GetLine(self, 0);
775 if (rest == NULL) {
776 Py_DECREF(line);
777 goto error;
778 }
779 PyString_Concat(&line, rest);
780 Py_DECREF(rest);
781 if (line == NULL)
782 goto error;
783 }
784 err = PyList_Append(list, line);
785 Py_DECREF(line);
786 if (err != 0)
787 goto error;
788 }
789
790 cleanup:
791 RELEASE_LOCK(self);
792 if (big_buffer) {
793 Py_DECREF(big_buffer);
794 }
795 return list;
796 }
797
798 PyDoc_STRVAR(BZ2File_xreadlines__doc__,
799 "xreadlines() -> self\n\
800 \n\
801 For backward compatibility. BZ2File objects now include the performance\n\
802 optimizations previously implemented in the xreadlines module.\n\
803 ");
804
805 PyDoc_STRVAR(BZ2File_write__doc__,
806 "write(data) -> None\n\
807 \n\
808 Write the 'data' string to file. Note that due to buffering, close() may\n\
809 be needed before the file on disk reflects the data written.\n\
810 ");
811
812 /* This is a hacked version of Python's fileobject.c:file_write(). */
813 static PyObject *
BZ2File_write(BZ2FileObject * self,PyObject * args)814 BZ2File_write(BZ2FileObject *self, PyObject *args)
815 {
816 PyObject *ret = NULL;
817 Py_buffer pbuf;
818 char *buf;
819 int len;
820 int bzerror;
821
822 if (!PyArg_ParseTuple(args, "s*:write", &pbuf))
823 return NULL;
824 buf = pbuf.buf;
825 len = pbuf.len;
826
827 ACQUIRE_LOCK(self);
828 switch (self->mode) {
829 case MODE_WRITE:
830 break;
831
832 case MODE_CLOSED:
833 PyErr_SetString(PyExc_ValueError,
834 "I/O operation on closed file");
835 goto cleanup;
836
837 default:
838 PyErr_SetString(PyExc_IOError,
839 "file is not ready for writing");
840 goto cleanup;
841 }
842
843 self->f_softspace = 0;
844
845 Py_BEGIN_ALLOW_THREADS
846 BZ2_bzWrite (&bzerror, self->fp, buf, len);
847 self->pos += len;
848 Py_END_ALLOW_THREADS
849
850 if (bzerror != BZ_OK) {
851 Util_CatchBZ2Error(bzerror);
852 goto cleanup;
853 }
854
855 Py_INCREF(Py_None);
856 ret = Py_None;
857
858 cleanup:
859 PyBuffer_Release(&pbuf);
860 RELEASE_LOCK(self);
861 return ret;
862 }
863
864 PyDoc_STRVAR(BZ2File_writelines__doc__,
865 "writelines(sequence_of_strings) -> None\n\
866 \n\
867 Write the sequence of strings to the file. Note that newlines are not\n\
868 added. The sequence can be any iterable object producing strings. This is\n\
869 equivalent to calling write() for each string.\n\
870 ");
871
872 /* This is a hacked version of Python's fileobject.c:file_writelines(). */
873 static PyObject *
BZ2File_writelines(BZ2FileObject * self,PyObject * seq)874 BZ2File_writelines(BZ2FileObject *self, PyObject *seq)
875 {
876 #define CHUNKSIZE 1000
877 PyObject *list = NULL;
878 PyObject *iter = NULL;
879 PyObject *ret = NULL;
880 PyObject *line;
881 int i, j, index, len, islist;
882 int bzerror;
883
884 ACQUIRE_LOCK(self);
885 switch (self->mode) {
886 case MODE_WRITE:
887 break;
888
889 case MODE_CLOSED:
890 PyErr_SetString(PyExc_ValueError,
891 "I/O operation on closed file");
892 goto error;
893
894 default:
895 PyErr_SetString(PyExc_IOError,
896 "file is not ready for writing");
897 goto error;
898 }
899
900 islist = PyList_Check(seq);
901 if (!islist) {
902 iter = PyObject_GetIter(seq);
903 if (iter == NULL) {
904 PyErr_SetString(PyExc_TypeError,
905 "writelines() requires an iterable argument");
906 goto error;
907 }
908 list = PyList_New(CHUNKSIZE);
909 if (list == NULL)
910 goto error;
911 }
912
913 /* Strategy: slurp CHUNKSIZE lines into a private list,
914 checking that they are all strings, then write that list
915 without holding the interpreter lock, then come back for more. */
916 for (index = 0; ; index += CHUNKSIZE) {
917 if (islist) {
918 Py_XDECREF(list);
919 list = PyList_GetSlice(seq, index, index+CHUNKSIZE);
920 if (list == NULL)
921 goto error;
922 j = PyList_GET_SIZE(list);
923 }
924 else {
925 for (j = 0; j < CHUNKSIZE; j++) {
926 line = PyIter_Next(iter);
927 if (line == NULL) {
928 if (PyErr_Occurred())
929 goto error;
930 break;
931 }
932 PyList_SetItem(list, j, line);
933 }
934 }
935 if (j == 0)
936 break;
937
938 /* Check that all entries are indeed strings. If not,
939 apply the same rules as for file.write() and
940 convert the rets to strings. This is slow, but
941 seems to be the only way since all conversion APIs
942 could potentially execute Python code. */
943 for (i = 0; i < j; i++) {
944 PyObject *v = PyList_GET_ITEM(list, i);
945 if (!PyString_Check(v)) {
946 const char *buffer;
947 Py_ssize_t len;
948 if (PyObject_AsCharBuffer(v, &buffer, &len)) {
949 PyErr_SetString(PyExc_TypeError,
950 "writelines() "
951 "argument must be "
952 "a sequence of "
953 "strings");
954 goto error;
955 }
956 line = PyString_FromStringAndSize(buffer,
957 len);
958 if (line == NULL)
959 goto error;
960 Py_DECREF(v);
961 PyList_SET_ITEM(list, i, line);
962 }
963 }
964
965 self->f_softspace = 0;
966
967 /* Since we are releasing the global lock, the
968 following code may *not* execute Python code. */
969 Py_BEGIN_ALLOW_THREADS
970 for (i = 0; i < j; i++) {
971 line = PyList_GET_ITEM(list, i);
972 len = PyString_GET_SIZE(line);
973 BZ2_bzWrite (&bzerror, self->fp,
974 PyString_AS_STRING(line), len);
975 if (bzerror != BZ_OK) {
976 Py_BLOCK_THREADS
977 Util_CatchBZ2Error(bzerror);
978 goto error;
979 }
980 }
981 Py_END_ALLOW_THREADS
982
983 if (j < CHUNKSIZE)
984 break;
985 }
986
987 Py_INCREF(Py_None);
988 ret = Py_None;
989
990 error:
991 RELEASE_LOCK(self);
992 Py_XDECREF(list);
993 Py_XDECREF(iter);
994 return ret;
995 #undef CHUNKSIZE
996 }
997
998 PyDoc_STRVAR(BZ2File_seek__doc__,
999 "seek(offset [, whence]) -> None\n\
1000 \n\
1001 Move to new file position. Argument offset is a byte count. Optional\n\
1002 argument whence defaults to 0 (offset from start of file, offset\n\
1003 should be >= 0); other values are 1 (move relative to current position,\n\
1004 positive or negative), and 2 (move relative to end of file, usually\n\
1005 negative, although many platforms allow seeking beyond the end of a file).\n\
1006 \n\
1007 Note that seeking of bz2 files is emulated, and depending on the parameters\n\
1008 the operation may be extremely slow.\n\
1009 ");
1010
1011 static PyObject *
BZ2File_seek(BZ2FileObject * self,PyObject * args)1012 BZ2File_seek(BZ2FileObject *self, PyObject *args)
1013 {
1014 int where = 0;
1015 PyObject *offobj;
1016 Py_off_t offset;
1017 char small_buffer[SMALLCHUNK];
1018 char *buffer = small_buffer;
1019 size_t buffersize = SMALLCHUNK;
1020 Py_off_t bytesread = 0;
1021 size_t readsize;
1022 int chunksize;
1023 int bzerror;
1024 PyObject *ret = NULL;
1025
1026 if (!PyArg_ParseTuple(args, "O|i:seek", &offobj, &where))
1027 return NULL;
1028 #if !defined(HAVE_LARGEFILE_SUPPORT)
1029 offset = PyInt_AsLong(offobj);
1030 #else
1031 offset = PyLong_Check(offobj) ?
1032 PyLong_AsLongLong(offobj) : PyInt_AsLong(offobj);
1033 #endif
1034 if (PyErr_Occurred())
1035 return NULL;
1036
1037 ACQUIRE_LOCK(self);
1038 Util_DropReadAhead(self);
1039 switch (self->mode) {
1040 case MODE_READ:
1041 case MODE_READ_EOF:
1042 break;
1043
1044 case MODE_CLOSED:
1045 PyErr_SetString(PyExc_ValueError,
1046 "I/O operation on closed file");
1047 goto cleanup;
1048
1049 default:
1050 PyErr_SetString(PyExc_IOError,
1051 "seek works only while reading");
1052 goto cleanup;
1053 }
1054
1055 if (where == 2) {
1056 if (self->size == -1) {
1057 assert(self->mode != MODE_READ_EOF);
1058 for (;;) {
1059 Py_BEGIN_ALLOW_THREADS
1060 chunksize = Util_UnivNewlineRead(
1061 &bzerror, self->fp,
1062 buffer, buffersize,
1063 self);
1064 self->pos += chunksize;
1065 Py_END_ALLOW_THREADS
1066
1067 bytesread += chunksize;
1068 if (bzerror == BZ_STREAM_END) {
1069 break;
1070 } else if (bzerror != BZ_OK) {
1071 Util_CatchBZ2Error(bzerror);
1072 goto cleanup;
1073 }
1074 }
1075 self->mode = MODE_READ_EOF;
1076 self->size = self->pos;
1077 bytesread = 0;
1078 }
1079 offset = self->size + offset;
1080 } else if (where == 1) {
1081 offset = self->pos + offset;
1082 }
1083
1084 /* Before getting here, offset must be the absolute position the file
1085 * pointer should be set to. */
1086
1087 if (offset >= self->pos) {
1088 /* we can move forward */
1089 offset -= self->pos;
1090 } else {
1091 /* we cannot move back, so rewind the stream */
1092 BZ2_bzReadClose(&bzerror, self->fp);
1093 if (self->fp) {
1094 PyFile_DecUseCount((PyFileObject *)self->file);
1095 self->fp = NULL;
1096 }
1097 if (bzerror != BZ_OK) {
1098 Util_CatchBZ2Error(bzerror);
1099 goto cleanup;
1100 }
1101 ret = PyObject_CallMethod(self->file, "seek", "(i)", 0);
1102 if (!ret)
1103 goto cleanup;
1104 Py_DECREF(ret);
1105 ret = NULL;
1106 self->pos = 0;
1107 self->fp = BZ2_bzReadOpen(&bzerror, PyFile_AsFile(self->file),
1108 0, 0, NULL, 0);
1109 if (self->fp)
1110 PyFile_IncUseCount((PyFileObject *)self->file);
1111 if (bzerror != BZ_OK) {
1112 Util_CatchBZ2Error(bzerror);
1113 goto cleanup;
1114 }
1115 self->mode = MODE_READ;
1116 }
1117
1118 if (offset <= 0 || self->mode == MODE_READ_EOF)
1119 goto exit;
1120
1121 /* Before getting here, offset must be set to the number of bytes
1122 * to walk forward. */
1123 for (;;) {
1124 if (offset-bytesread > buffersize)
1125 readsize = buffersize;
1126 else
1127 /* offset might be wider that readsize, but the result
1128 * of the subtraction is bound by buffersize (see the
1129 * condition above). buffersize is 8192. */
1130 readsize = (size_t)(offset-bytesread);
1131 Py_BEGIN_ALLOW_THREADS
1132 chunksize = Util_UnivNewlineRead(&bzerror, self->fp,
1133 buffer, readsize, self);
1134 self->pos += chunksize;
1135 Py_END_ALLOW_THREADS
1136 bytesread += chunksize;
1137 if (bzerror == BZ_STREAM_END) {
1138 self->size = self->pos;
1139 self->mode = MODE_READ_EOF;
1140 break;
1141 } else if (bzerror != BZ_OK) {
1142 Util_CatchBZ2Error(bzerror);
1143 goto cleanup;
1144 }
1145 if (bytesread == offset)
1146 break;
1147 }
1148
1149 exit:
1150 Py_INCREF(Py_None);
1151 ret = Py_None;
1152
1153 cleanup:
1154 RELEASE_LOCK(self);
1155 return ret;
1156 }
1157
1158 PyDoc_STRVAR(BZ2File_tell__doc__,
1159 "tell() -> int\n\
1160 \n\
1161 Return the current file position, an integer (may be a long integer).\n\
1162 ");
1163
1164 static PyObject *
BZ2File_tell(BZ2FileObject * self,PyObject * args)1165 BZ2File_tell(BZ2FileObject *self, PyObject *args)
1166 {
1167 PyObject *ret = NULL;
1168
1169 if (self->mode == MODE_CLOSED) {
1170 PyErr_SetString(PyExc_ValueError,
1171 "I/O operation on closed file");
1172 goto cleanup;
1173 }
1174
1175 #if !defined(HAVE_LARGEFILE_SUPPORT)
1176 ret = PyInt_FromLong(self->pos);
1177 #else
1178 ret = PyLong_FromLongLong(self->pos);
1179 #endif
1180
1181 cleanup:
1182 return ret;
1183 }
1184
1185 PyDoc_STRVAR(BZ2File_close__doc__,
1186 "close() -> None or (perhaps) an integer\n\
1187 \n\
1188 Close the file. Sets data attribute .closed to true. A closed file\n\
1189 cannot be used for further I/O operations. close() may be called more\n\
1190 than once without error.\n\
1191 ");
1192
1193 static PyObject *
BZ2File_close(BZ2FileObject * self)1194 BZ2File_close(BZ2FileObject *self)
1195 {
1196 PyObject *ret = NULL;
1197 int bzerror = BZ_OK;
1198
1199 ACQUIRE_LOCK(self);
1200 switch (self->mode) {
1201 case MODE_READ:
1202 case MODE_READ_EOF:
1203 BZ2_bzReadClose(&bzerror, self->fp);
1204 break;
1205 case MODE_WRITE:
1206 BZ2_bzWriteClose(&bzerror, self->fp,
1207 0, NULL, NULL);
1208 break;
1209 }
1210 if (self->file) {
1211 if (self->fp)
1212 PyFile_DecUseCount((PyFileObject *)self->file);
1213 ret = PyObject_CallMethod(self->file, "close", NULL);
1214 } else {
1215 Py_INCREF(Py_None);
1216 ret = Py_None;
1217 }
1218 self->fp = NULL;
1219 self->mode = MODE_CLOSED;
1220 if (bzerror != BZ_OK) {
1221 Util_CatchBZ2Error(bzerror);
1222 Py_XDECREF(ret);
1223 ret = NULL;
1224 }
1225
1226 RELEASE_LOCK(self);
1227 return ret;
1228 }
1229
1230 PyDoc_STRVAR(BZ2File_enter_doc,
1231 "__enter__() -> self.");
1232
1233 static PyObject *
BZ2File_enter(BZ2FileObject * self)1234 BZ2File_enter(BZ2FileObject *self)
1235 {
1236 if (self->mode == MODE_CLOSED) {
1237 PyErr_SetString(PyExc_ValueError,
1238 "I/O operation on closed file");
1239 return NULL;
1240 }
1241 Py_INCREF(self);
1242 return (PyObject *) self;
1243 }
1244
1245 PyDoc_STRVAR(BZ2File_exit_doc,
1246 "__exit__(*excinfo) -> None. Closes the file.");
1247
1248 static PyObject *
BZ2File_exit(BZ2FileObject * self,PyObject * args)1249 BZ2File_exit(BZ2FileObject *self, PyObject *args)
1250 {
1251 PyObject *ret = PyObject_CallMethod((PyObject *) self, "close", NULL);
1252 if (!ret)
1253 /* If error occurred, pass through */
1254 return NULL;
1255 Py_DECREF(ret);
1256 Py_RETURN_NONE;
1257 }
1258
1259
1260 static PyObject *BZ2File_getiter(BZ2FileObject *self);
1261
1262 static PyMethodDef BZ2File_methods[] = {
1263 {"read", (PyCFunction)BZ2File_read, METH_VARARGS, BZ2File_read__doc__},
1264 {"readline", (PyCFunction)BZ2File_readline, METH_VARARGS, BZ2File_readline__doc__},
1265 {"readlines", (PyCFunction)BZ2File_readlines, METH_VARARGS, BZ2File_readlines__doc__},
1266 {"xreadlines", (PyCFunction)BZ2File_getiter, METH_VARARGS, BZ2File_xreadlines__doc__},
1267 {"write", (PyCFunction)BZ2File_write, METH_VARARGS, BZ2File_write__doc__},
1268 {"writelines", (PyCFunction)BZ2File_writelines, METH_O, BZ2File_writelines__doc__},
1269 {"seek", (PyCFunction)BZ2File_seek, METH_VARARGS, BZ2File_seek__doc__},
1270 {"tell", (PyCFunction)BZ2File_tell, METH_NOARGS, BZ2File_tell__doc__},
1271 {"close", (PyCFunction)BZ2File_close, METH_NOARGS, BZ2File_close__doc__},
1272 {"__enter__", (PyCFunction)BZ2File_enter, METH_NOARGS, BZ2File_enter_doc},
1273 {"__exit__", (PyCFunction)BZ2File_exit, METH_VARARGS, BZ2File_exit_doc},
1274 {NULL, NULL} /* sentinel */
1275 };
1276
1277
1278 /* ===================================================================== */
1279 /* Getters and setters of BZ2File. */
1280
1281 /* This is a hacked version of Python's fileobject.c:get_newlines(). */
1282 static PyObject *
BZ2File_get_newlines(BZ2FileObject * self,void * closure)1283 BZ2File_get_newlines(BZ2FileObject *self, void *closure)
1284 {
1285 switch (self->f_newlinetypes) {
1286 case NEWLINE_UNKNOWN:
1287 Py_INCREF(Py_None);
1288 return Py_None;
1289 case NEWLINE_CR:
1290 return PyString_FromString("\r");
1291 case NEWLINE_LF:
1292 return PyString_FromString("\n");
1293 case NEWLINE_CR|NEWLINE_LF:
1294 return Py_BuildValue("(ss)", "\r", "\n");
1295 case NEWLINE_CRLF:
1296 return PyString_FromString("\r\n");
1297 case NEWLINE_CR|NEWLINE_CRLF:
1298 return Py_BuildValue("(ss)", "\r", "\r\n");
1299 case NEWLINE_LF|NEWLINE_CRLF:
1300 return Py_BuildValue("(ss)", "\n", "\r\n");
1301 case NEWLINE_CR|NEWLINE_LF|NEWLINE_CRLF:
1302 return Py_BuildValue("(sss)", "\r", "\n", "\r\n");
1303 default:
1304 PyErr_Format(PyExc_SystemError,
1305 "Unknown newlines value 0x%x\n",
1306 self->f_newlinetypes);
1307 return NULL;
1308 }
1309 }
1310
1311 static PyObject *
BZ2File_get_closed(BZ2FileObject * self,void * closure)1312 BZ2File_get_closed(BZ2FileObject *self, void *closure)
1313 {
1314 return PyInt_FromLong(self->mode == MODE_CLOSED);
1315 }
1316
1317 static PyObject *
BZ2File_get_mode(BZ2FileObject * self,void * closure)1318 BZ2File_get_mode(BZ2FileObject *self, void *closure)
1319 {
1320 return PyObject_GetAttrString(self->file, "mode");
1321 }
1322
1323 static PyObject *
BZ2File_get_name(BZ2FileObject * self,void * closure)1324 BZ2File_get_name(BZ2FileObject *self, void *closure)
1325 {
1326 return PyObject_GetAttrString(self->file, "name");
1327 }
1328
1329 static PyGetSetDef BZ2File_getset[] = {
1330 {"closed", (getter)BZ2File_get_closed, NULL,
1331 "True if the file is closed"},
1332 {"newlines", (getter)BZ2File_get_newlines, NULL,
1333 "end-of-line convention used in this file"},
1334 {"mode", (getter)BZ2File_get_mode, NULL,
1335 "file mode ('r', 'w', or 'U')"},
1336 {"name", (getter)BZ2File_get_name, NULL,
1337 "file name"},
1338 {NULL} /* Sentinel */
1339 };
1340
1341
1342 /* ===================================================================== */
1343 /* Members of BZ2File_Type. */
1344
1345 #undef OFF
1346 #define OFF(x) offsetof(BZ2FileObject, x)
1347
1348 static PyMemberDef BZ2File_members[] = {
1349 {"softspace", T_INT, OFF(f_softspace), 0,
1350 "flag indicating that a space needs to be printed; used by print"},
1351 {NULL} /* Sentinel */
1352 };
1353
1354 /* ===================================================================== */
1355 /* Slot definitions for BZ2File_Type. */
1356 static int
BZ2File_clear(BZ2FileObject * self)1357 BZ2File_clear(BZ2FileObject *self)
1358 {
1359 int bzerror;
1360
1361 ACQUIRE_LOCK(self);
1362 switch (self->mode) {
1363 case MODE_READ:
1364 case MODE_READ_EOF:
1365 BZ2_bzReadClose(&bzerror, self->fp);
1366 break;
1367 case MODE_WRITE:
1368 BZ2_bzWriteClose(&bzerror, self->fp,
1369 0, NULL, NULL);
1370 break;
1371 }
1372 if (self->fp != NULL && self->file != NULL)
1373 PyFile_DecUseCount((PyFileObject *)self->file);
1374 self->fp = NULL;
1375 Util_DropReadAhead(self);
1376 Py_CLEAR(self->file);
1377 RELEASE_LOCK(self);
1378 return 0;
1379 }
1380
1381 static int
BZ2File_init(BZ2FileObject * self,PyObject * args,PyObject * kwargs)1382 BZ2File_init(BZ2FileObject *self, PyObject *args, PyObject *kwargs)
1383 {
1384 static char *kwlist[] = {"filename", "mode", "buffering",
1385 "compresslevel", 0};
1386 PyObject *name;
1387 char *mode = "r";
1388 int buffering = -1;
1389 int compresslevel = 9;
1390 int bzerror;
1391 int mode_char = 0;
1392
1393 self->size = -1;
1394
1395 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O|sii:BZ2File",
1396 kwlist, &name, &mode, &buffering,
1397 &compresslevel))
1398 return -1;
1399
1400 if (compresslevel < 1 || compresslevel > 9) {
1401 PyErr_SetString(PyExc_ValueError,
1402 "compresslevel must be between 1 and 9");
1403 return -1;
1404 }
1405
1406 for (;;) {
1407 int error = 0;
1408 switch (*mode) {
1409 case 'r':
1410 case 'w':
1411 if (mode_char)
1412 error = 1;
1413 mode_char = *mode;
1414 break;
1415
1416 case 'b':
1417 break;
1418
1419 case 'U':
1420 #ifdef __VMS
1421 self->f_univ_newline = 0;
1422 #else
1423 self->f_univ_newline = 1;
1424 #endif
1425 break;
1426
1427 default:
1428 error = 1;
1429 break;
1430 }
1431 if (error) {
1432 PyErr_Format(PyExc_ValueError,
1433 "invalid mode char %c", *mode);
1434 return -1;
1435 }
1436 mode++;
1437 if (*mode == '\0')
1438 break;
1439 }
1440
1441 if (mode_char == 0) {
1442 mode_char = 'r';
1443 }
1444
1445 mode = (mode_char == 'r') ? "rb" : "wb";
1446
1447 #ifdef WITH_THREAD
1448 if (!self->lock) {
1449 self->lock = PyThread_allocate_lock();
1450 }
1451
1452 if (!self->lock) {
1453 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1454 goto error;
1455 }
1456 #endif
1457
1458 BZ2File_clear(self);
1459
1460 self->file = PyObject_CallFunction((PyObject*)&PyFile_Type, "(Osi)",
1461 name, mode, buffering);
1462 if (self->file == NULL)
1463 return -1;
1464
1465 /* From now on, we have stuff to dealloc, so jump to error label
1466 * instead of returning */
1467
1468 if (mode_char == 'r')
1469 self->fp = BZ2_bzReadOpen(&bzerror,
1470 PyFile_AsFile(self->file),
1471 0, 0, NULL, 0);
1472 else
1473 self->fp = BZ2_bzWriteOpen(&bzerror,
1474 PyFile_AsFile(self->file),
1475 compresslevel, 0, 0);
1476
1477 if (bzerror != BZ_OK) {
1478 Util_CatchBZ2Error(bzerror);
1479 goto error;
1480 }
1481 PyFile_IncUseCount((PyFileObject *)self->file);
1482
1483 self->mode = (mode_char == 'r') ? MODE_READ : MODE_WRITE;
1484
1485 return 0;
1486
1487 error:
1488 Py_CLEAR(self->file);
1489 #ifdef WITH_THREAD
1490 if (self->lock) {
1491 PyThread_free_lock(self->lock);
1492 self->lock = NULL;
1493 }
1494 #endif
1495 return -1;
1496 }
1497
1498 static void
BZ2File_dealloc(BZ2FileObject * self)1499 BZ2File_dealloc(BZ2FileObject *self)
1500 {
1501 BZ2File_clear(self);
1502 #ifdef WITH_THREAD
1503 if (self->lock)
1504 PyThread_free_lock(self->lock);
1505 #endif
1506 Py_TYPE(self)->tp_free((PyObject *)self);
1507 }
1508
1509 /* This is a hacked version of Python's fileobject.c:file_getiter(). */
1510 static PyObject *
BZ2File_getiter(BZ2FileObject * self)1511 BZ2File_getiter(BZ2FileObject *self)
1512 {
1513 if (self->mode == MODE_CLOSED) {
1514 PyErr_SetString(PyExc_ValueError,
1515 "I/O operation on closed file");
1516 return NULL;
1517 }
1518 Py_INCREF((PyObject*)self);
1519 return (PyObject *)self;
1520 }
1521
1522 /* This is a hacked version of Python's fileobject.c:file_iternext(). */
1523 #define READAHEAD_BUFSIZE 8192
1524 static PyObject *
BZ2File_iternext(BZ2FileObject * self)1525 BZ2File_iternext(BZ2FileObject *self)
1526 {
1527 PyStringObject* ret;
1528 ACQUIRE_LOCK(self);
1529 if (self->mode == MODE_CLOSED) {
1530 RELEASE_LOCK(self);
1531 PyErr_SetString(PyExc_ValueError,
1532 "I/O operation on closed file");
1533 return NULL;
1534 }
1535 ret = Util_ReadAheadGetLineSkip(self, 0, READAHEAD_BUFSIZE);
1536 RELEASE_LOCK(self);
1537 if (ret == NULL || PyString_GET_SIZE(ret) == 0) {
1538 Py_XDECREF(ret);
1539 return NULL;
1540 }
1541 return (PyObject *)ret;
1542 }
1543
1544 /* ===================================================================== */
1545 /* BZ2File_Type definition. */
1546
1547 PyDoc_VAR(BZ2File__doc__) =
1548 PyDoc_STR(
1549 "BZ2File(name [, mode='r', buffering=0, compresslevel=9]) -> file object\n\
1550 \n\
1551 Open a bz2 file. The mode can be 'r' or 'w', for reading (default) or\n\
1552 writing. When opened for writing, the file will be created if it doesn't\n\
1553 exist, and truncated otherwise. If the buffering argument is given, 0 means\n\
1554 unbuffered, and larger numbers specify the buffer size. If compresslevel\n\
1555 is given, must be a number between 1 and 9.\n\
1556 ")
1557 PyDoc_STR(
1558 "\n\
1559 Add a 'U' to mode to open the file for input with universal newline\n\
1560 support. Any line ending in the input file will be seen as a '\\n' in\n\
1561 Python. Also, a file so opened gains the attribute 'newlines'; the value\n\
1562 for this attribute is one of None (no newline read yet), '\\r', '\\n',\n\
1563 '\\r\\n' or a tuple containing all the newline types seen. Universal\n\
1564 newlines are available only when reading.\n\
1565 ")
1566 ;
1567
1568 static PyTypeObject BZ2File_Type = {
1569 PyVarObject_HEAD_INIT(NULL, 0)
1570 "bz2.BZ2File", /*tp_name*/
1571 sizeof(BZ2FileObject), /*tp_basicsize*/
1572 0, /*tp_itemsize*/
1573 (destructor)BZ2File_dealloc, /*tp_dealloc*/
1574 0, /*tp_print*/
1575 0, /*tp_getattr*/
1576 0, /*tp_setattr*/
1577 0, /*tp_compare*/
1578 0, /*tp_repr*/
1579 0, /*tp_as_number*/
1580 0, /*tp_as_sequence*/
1581 0, /*tp_as_mapping*/
1582 0, /*tp_hash*/
1583 0, /*tp_call*/
1584 0, /*tp_str*/
1585 PyObject_GenericGetAttr,/*tp_getattro*/
1586 PyObject_GenericSetAttr,/*tp_setattro*/
1587 0, /*tp_as_buffer*/
1588 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1589 BZ2File__doc__, /*tp_doc*/
1590 0, /*tp_traverse*/
1591 (inquiry)BZ2File_clear, /*tp_clear*/
1592 0, /*tp_richcompare*/
1593 0, /*tp_weaklistoffset*/
1594 (getiterfunc)BZ2File_getiter, /*tp_iter*/
1595 (iternextfunc)BZ2File_iternext, /*tp_iternext*/
1596 BZ2File_methods, /*tp_methods*/
1597 BZ2File_members, /*tp_members*/
1598 BZ2File_getset, /*tp_getset*/
1599 0, /*tp_base*/
1600 0, /*tp_dict*/
1601 0, /*tp_descr_get*/
1602 0, /*tp_descr_set*/
1603 0, /*tp_dictoffset*/
1604 (initproc)BZ2File_init, /*tp_init*/
1605 PyType_GenericAlloc, /*tp_alloc*/
1606 PyType_GenericNew, /*tp_new*/
1607 _PyObject_Del, /*tp_free*/
1608 0, /*tp_is_gc*/
1609 };
1610
1611
1612 /* ===================================================================== */
1613 /* Methods of BZ2Comp. */
1614
1615 PyDoc_STRVAR(BZ2Comp_compress__doc__,
1616 "compress(data) -> string\n\
1617 \n\
1618 Provide more data to the compressor object. It will return chunks of\n\
1619 compressed data whenever possible. When you've finished providing data\n\
1620 to compress, call the flush() method to finish the compression process,\n\
1621 and return what is left in the internal buffers.\n\
1622 ");
1623
1624 static PyObject *
BZ2Comp_compress(BZ2CompObject * self,PyObject * args)1625 BZ2Comp_compress(BZ2CompObject *self, PyObject *args)
1626 {
1627 Py_buffer pdata;
1628 size_t input_left;
1629 size_t output_size = 0;
1630 PyObject *ret = NULL;
1631 bz_stream *bzs = &self->bzs;
1632 int bzerror;
1633
1634 if (!PyArg_ParseTuple(args, "s*:compress", &pdata))
1635 return NULL;
1636
1637 if (pdata.len == 0) {
1638 PyBuffer_Release(&pdata);
1639 return PyString_FromString("");
1640 }
1641
1642 ACQUIRE_LOCK(self);
1643 if (!self->running) {
1644 PyErr_SetString(PyExc_ValueError,
1645 "this object was already flushed");
1646 goto error;
1647 }
1648
1649 ret = PyString_FromStringAndSize(NULL, SMALLCHUNK);
1650 if (!ret)
1651 goto error;
1652
1653 bzs->next_in = pdata.buf;
1654 bzs->avail_in = MIN(pdata.len, UINT_MAX);
1655 input_left = pdata.len - bzs->avail_in;
1656
1657 bzs->next_out = BUF(ret);
1658 bzs->avail_out = PyString_GET_SIZE(ret);
1659
1660 for (;;) {
1661 char *saved_next_out;
1662
1663 Py_BEGIN_ALLOW_THREADS
1664 saved_next_out = bzs->next_out;
1665 bzerror = BZ2_bzCompress(bzs, BZ_RUN);
1666 output_size += bzs->next_out - saved_next_out;
1667 Py_END_ALLOW_THREADS
1668
1669 if (bzerror != BZ_RUN_OK) {
1670 Util_CatchBZ2Error(bzerror);
1671 goto error;
1672 }
1673 if (bzs->avail_in == 0) {
1674 if (input_left == 0)
1675 break; /* no more input data */
1676 bzs->avail_in = MIN(input_left, UINT_MAX);
1677 input_left -= bzs->avail_in;
1678 }
1679 if (bzs->avail_out == 0) {
1680 size_t buffer_left = PyString_GET_SIZE(ret) - output_size;
1681 if (buffer_left == 0) {
1682 if (Util_GrowBuffer(&ret) < 0) {
1683 BZ2_bzCompressEnd(bzs);
1684 goto error;
1685 }
1686 bzs->next_out = BUF(ret) + output_size;
1687 buffer_left = PyString_GET_SIZE(ret) - output_size;
1688 }
1689 bzs->avail_out = MIN(buffer_left, UINT_MAX);
1690 }
1691 }
1692
1693 if (_PyString_Resize(&ret, output_size) < 0)
1694 goto error;
1695
1696 RELEASE_LOCK(self);
1697 PyBuffer_Release(&pdata);
1698 return ret;
1699
1700 error:
1701 RELEASE_LOCK(self);
1702 PyBuffer_Release(&pdata);
1703 Py_XDECREF(ret);
1704 return NULL;
1705 }
1706
1707 PyDoc_STRVAR(BZ2Comp_flush__doc__,
1708 "flush() -> string\n\
1709 \n\
1710 Finish the compression process and return what is left in internal buffers.\n\
1711 You must not use the compressor object after calling this method.\n\
1712 ");
1713
1714 static PyObject *
BZ2Comp_flush(BZ2CompObject * self)1715 BZ2Comp_flush(BZ2CompObject *self)
1716 {
1717 size_t output_size = 0;
1718 PyObject *ret = NULL;
1719 bz_stream *bzs = &self->bzs;
1720 int bzerror;
1721
1722 ACQUIRE_LOCK(self);
1723 if (!self->running) {
1724 PyErr_SetString(PyExc_ValueError, "object was already flushed");
1725 goto error;
1726 }
1727 self->running = 0;
1728
1729 ret = PyString_FromStringAndSize(NULL, SMALLCHUNK);
1730 if (!ret)
1731 goto error;
1732
1733 bzs->next_out = BUF(ret);
1734 bzs->avail_out = PyString_GET_SIZE(ret);
1735
1736 for (;;) {
1737 char *saved_next_out;
1738
1739 Py_BEGIN_ALLOW_THREADS
1740 saved_next_out = bzs->next_out;
1741 bzerror = BZ2_bzCompress(bzs, BZ_FINISH);
1742 output_size += bzs->next_out - saved_next_out;
1743 Py_END_ALLOW_THREADS
1744
1745 if (bzerror == BZ_STREAM_END) {
1746 break;
1747 } else if (bzerror != BZ_FINISH_OK) {
1748 Util_CatchBZ2Error(bzerror);
1749 goto error;
1750 }
1751 if (bzs->avail_out == 0) {
1752 size_t buffer_left = PyString_GET_SIZE(ret) - output_size;
1753 if (buffer_left == 0) {
1754 if (Util_GrowBuffer(&ret) < 0)
1755 goto error;
1756 bzs->next_out = BUF(ret) + output_size;
1757 buffer_left = PyString_GET_SIZE(ret) - output_size;
1758 }
1759 bzs->avail_out = MIN(buffer_left, UINT_MAX);
1760 }
1761 }
1762
1763 if (output_size != PyString_GET_SIZE(ret))
1764 if (_PyString_Resize(&ret, output_size) < 0)
1765 goto error;
1766
1767 RELEASE_LOCK(self);
1768 return ret;
1769
1770 error:
1771 RELEASE_LOCK(self);
1772 Py_XDECREF(ret);
1773 return NULL;
1774 }
1775
1776 static PyMethodDef BZ2Comp_methods[] = {
1777 {"compress", (PyCFunction)BZ2Comp_compress, METH_VARARGS,
1778 BZ2Comp_compress__doc__},
1779 {"flush", (PyCFunction)BZ2Comp_flush, METH_NOARGS,
1780 BZ2Comp_flush__doc__},
1781 {NULL, NULL} /* sentinel */
1782 };
1783
1784
1785 /* ===================================================================== */
1786 /* Slot definitions for BZ2Comp_Type. */
1787
1788 static int
BZ2Comp_init(BZ2CompObject * self,PyObject * args,PyObject * kwargs)1789 BZ2Comp_init(BZ2CompObject *self, PyObject *args, PyObject *kwargs)
1790 {
1791 int compresslevel = 9;
1792 int bzerror;
1793 static char *kwlist[] = {"compresslevel", 0};
1794
1795 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "|i:BZ2Compressor",
1796 kwlist, &compresslevel))
1797 return -1;
1798
1799 if (compresslevel < 1 || compresslevel > 9) {
1800 PyErr_SetString(PyExc_ValueError,
1801 "compresslevel must be between 1 and 9");
1802 goto error;
1803 }
1804
1805 #ifdef WITH_THREAD
1806 self->lock = PyThread_allocate_lock();
1807 if (!self->lock) {
1808 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
1809 goto error;
1810 }
1811 #endif
1812
1813 memset(&self->bzs, 0, sizeof(bz_stream));
1814 bzerror = BZ2_bzCompressInit(&self->bzs, compresslevel, 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 error:
1824 #ifdef WITH_THREAD
1825 if (self->lock) {
1826 PyThread_free_lock(self->lock);
1827 self->lock = NULL;
1828 }
1829 #endif
1830 return -1;
1831 }
1832
1833 static void
BZ2Comp_dealloc(BZ2CompObject * self)1834 BZ2Comp_dealloc(BZ2CompObject *self)
1835 {
1836 #ifdef WITH_THREAD
1837 if (self->lock)
1838 PyThread_free_lock(self->lock);
1839 #endif
1840 BZ2_bzCompressEnd(&self->bzs);
1841 Py_TYPE(self)->tp_free((PyObject *)self);
1842 }
1843
1844
1845 /* ===================================================================== */
1846 /* BZ2Comp_Type definition. */
1847
1848 PyDoc_STRVAR(BZ2Comp__doc__,
1849 "BZ2Compressor([compresslevel=9]) -> compressor object\n\
1850 \n\
1851 Create a new compressor object. This object may be used to compress\n\
1852 data sequentially. If you want to compress data in one shot, use the\n\
1853 compress() function instead. The compresslevel parameter, if given,\n\
1854 must be a number between 1 and 9.\n\
1855 ");
1856
1857 static PyTypeObject BZ2Comp_Type = {
1858 PyVarObject_HEAD_INIT(NULL, 0)
1859 "bz2.BZ2Compressor", /*tp_name*/
1860 sizeof(BZ2CompObject), /*tp_basicsize*/
1861 0, /*tp_itemsize*/
1862 (destructor)BZ2Comp_dealloc, /*tp_dealloc*/
1863 0, /*tp_print*/
1864 0, /*tp_getattr*/
1865 0, /*tp_setattr*/
1866 0, /*tp_compare*/
1867 0, /*tp_repr*/
1868 0, /*tp_as_number*/
1869 0, /*tp_as_sequence*/
1870 0, /*tp_as_mapping*/
1871 0, /*tp_hash*/
1872 0, /*tp_call*/
1873 0, /*tp_str*/
1874 PyObject_GenericGetAttr,/*tp_getattro*/
1875 PyObject_GenericSetAttr,/*tp_setattro*/
1876 0, /*tp_as_buffer*/
1877 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
1878 BZ2Comp__doc__, /*tp_doc*/
1879 0, /*tp_traverse*/
1880 0, /*tp_clear*/
1881 0, /*tp_richcompare*/
1882 0, /*tp_weaklistoffset*/
1883 0, /*tp_iter*/
1884 0, /*tp_iternext*/
1885 BZ2Comp_methods, /*tp_methods*/
1886 0, /*tp_members*/
1887 0, /*tp_getset*/
1888 0, /*tp_base*/
1889 0, /*tp_dict*/
1890 0, /*tp_descr_get*/
1891 0, /*tp_descr_set*/
1892 0, /*tp_dictoffset*/
1893 (initproc)BZ2Comp_init, /*tp_init*/
1894 PyType_GenericAlloc, /*tp_alloc*/
1895 PyType_GenericNew, /*tp_new*/
1896 _PyObject_Del, /*tp_free*/
1897 0, /*tp_is_gc*/
1898 };
1899
1900
1901 /* ===================================================================== */
1902 /* Members of BZ2Decomp. */
1903
1904 #undef OFF
1905 #define OFF(x) offsetof(BZ2DecompObject, x)
1906
1907 static PyMemberDef BZ2Decomp_members[] = {
1908 {"unused_data", T_OBJECT, OFF(unused_data), RO},
1909 {NULL} /* Sentinel */
1910 };
1911
1912
1913 /* ===================================================================== */
1914 /* Methods of BZ2Decomp. */
1915
1916 PyDoc_STRVAR(BZ2Decomp_decompress__doc__,
1917 "decompress(data) -> string\n\
1918 \n\
1919 Provide more data to the decompressor object. It will return chunks\n\
1920 of decompressed data whenever possible. If you try to decompress data\n\
1921 after the end of stream is found, EOFError will be raised. If any data\n\
1922 was found after the end of stream, it'll be ignored and saved in\n\
1923 unused_data attribute.\n\
1924 ");
1925
1926 static PyObject *
BZ2Decomp_decompress(BZ2DecompObject * self,PyObject * args)1927 BZ2Decomp_decompress(BZ2DecompObject *self, PyObject *args)
1928 {
1929 Py_buffer pdata;
1930 size_t input_left;
1931 size_t output_size = 0;
1932 PyObject *ret = NULL;
1933 bz_stream *bzs = &self->bzs;
1934 int bzerror;
1935
1936 if (!PyArg_ParseTuple(args, "s*:decompress", &pdata))
1937 return NULL;
1938
1939 ACQUIRE_LOCK(self);
1940 if (!self->running) {
1941 PyErr_SetString(PyExc_EOFError, "end of stream was "
1942 "already found");
1943 goto error;
1944 }
1945
1946 ret = PyString_FromStringAndSize(NULL, SMALLCHUNK);
1947 if (!ret)
1948 goto error;
1949
1950 bzs->next_in = pdata.buf;
1951 bzs->avail_in = MIN(pdata.len, UINT_MAX);
1952 input_left = pdata.len - bzs->avail_in;
1953
1954 bzs->next_out = BUF(ret);
1955 bzs->avail_out = PyString_GET_SIZE(ret);
1956
1957 for (;;) {
1958 char *saved_next_out;
1959
1960 Py_BEGIN_ALLOW_THREADS
1961 saved_next_out = bzs->next_out;
1962 bzerror = BZ2_bzDecompress(bzs);
1963 output_size += bzs->next_out - saved_next_out;
1964 Py_END_ALLOW_THREADS
1965
1966 if (bzerror == BZ_STREAM_END) {
1967 self->running = 0;
1968 input_left += bzs->avail_in;
1969 if (input_left != 0) {
1970 Py_SETREF(self->unused_data,
1971 PyString_FromStringAndSize(bzs->next_in, input_left));
1972 if (self->unused_data == NULL)
1973 goto error;
1974 }
1975 break;
1976 }
1977 if (bzerror != BZ_OK) {
1978 Util_CatchBZ2Error(bzerror);
1979 goto error;
1980 }
1981 if (bzs->avail_in == 0) {
1982 if (input_left == 0)
1983 break; /* no more input data */
1984 bzs->avail_in = MIN(input_left, UINT_MAX);
1985 input_left -= bzs->avail_in;
1986 }
1987 if (bzs->avail_out == 0) {
1988 size_t buffer_left = PyString_GET_SIZE(ret) - output_size;
1989 if (buffer_left == 0) {
1990 if (Util_GrowBuffer(&ret) < 0) {
1991 BZ2_bzDecompressEnd(bzs);
1992 goto error;
1993 }
1994 bzs->next_out = BUF(ret) + output_size;
1995 buffer_left = PyString_GET_SIZE(ret) - output_size;
1996 }
1997 bzs->avail_out = MIN(buffer_left, UINT_MAX);
1998 }
1999 }
2000
2001 if (output_size != PyString_GET_SIZE(ret))
2002 if (_PyString_Resize(&ret, output_size) < 0)
2003 goto error;
2004
2005 RELEASE_LOCK(self);
2006 PyBuffer_Release(&pdata);
2007 return ret;
2008
2009 error:
2010 RELEASE_LOCK(self);
2011 PyBuffer_Release(&pdata);
2012 Py_XDECREF(ret);
2013 return NULL;
2014 }
2015
2016 static PyMethodDef BZ2Decomp_methods[] = {
2017 {"decompress", (PyCFunction)BZ2Decomp_decompress, METH_VARARGS, BZ2Decomp_decompress__doc__},
2018 {NULL, NULL} /* sentinel */
2019 };
2020
2021
2022 /* ===================================================================== */
2023 /* Slot definitions for BZ2Decomp_Type. */
2024
2025 static int
BZ2Decomp_init(BZ2DecompObject * self,PyObject * args,PyObject * kwargs)2026 BZ2Decomp_init(BZ2DecompObject *self, PyObject *args, PyObject *kwargs)
2027 {
2028 int bzerror;
2029
2030 if (!PyArg_ParseTuple(args, ":BZ2Decompressor"))
2031 return -1;
2032
2033 #ifdef WITH_THREAD
2034 self->lock = PyThread_allocate_lock();
2035 if (!self->lock) {
2036 PyErr_SetString(PyExc_MemoryError, "unable to allocate lock");
2037 goto error;
2038 }
2039 #endif
2040
2041 self->unused_data = PyString_FromString("");
2042 if (!self->unused_data)
2043 goto error;
2044
2045 memset(&self->bzs, 0, sizeof(bz_stream));
2046 bzerror = BZ2_bzDecompressInit(&self->bzs, 0, 0);
2047 if (bzerror != BZ_OK) {
2048 Util_CatchBZ2Error(bzerror);
2049 goto error;
2050 }
2051
2052 self->running = 1;
2053
2054 return 0;
2055
2056 error:
2057 #ifdef WITH_THREAD
2058 if (self->lock) {
2059 PyThread_free_lock(self->lock);
2060 self->lock = NULL;
2061 }
2062 #endif
2063 Py_CLEAR(self->unused_data);
2064 return -1;
2065 }
2066
2067 static void
BZ2Decomp_dealloc(BZ2DecompObject * self)2068 BZ2Decomp_dealloc(BZ2DecompObject *self)
2069 {
2070 #ifdef WITH_THREAD
2071 if (self->lock)
2072 PyThread_free_lock(self->lock);
2073 #endif
2074 Py_XDECREF(self->unused_data);
2075 BZ2_bzDecompressEnd(&self->bzs);
2076 Py_TYPE(self)->tp_free((PyObject *)self);
2077 }
2078
2079
2080 /* ===================================================================== */
2081 /* BZ2Decomp_Type definition. */
2082
2083 PyDoc_STRVAR(BZ2Decomp__doc__,
2084 "BZ2Decompressor() -> decompressor object\n\
2085 \n\
2086 Create a new decompressor object. This object may be used to decompress\n\
2087 data sequentially. If you want to decompress data in one shot, use the\n\
2088 decompress() function instead.\n\
2089 ");
2090
2091 static PyTypeObject BZ2Decomp_Type = {
2092 PyVarObject_HEAD_INIT(NULL, 0)
2093 "bz2.BZ2Decompressor", /*tp_name*/
2094 sizeof(BZ2DecompObject), /*tp_basicsize*/
2095 0, /*tp_itemsize*/
2096 (destructor)BZ2Decomp_dealloc, /*tp_dealloc*/
2097 0, /*tp_print*/
2098 0, /*tp_getattr*/
2099 0, /*tp_setattr*/
2100 0, /*tp_compare*/
2101 0, /*tp_repr*/
2102 0, /*tp_as_number*/
2103 0, /*tp_as_sequence*/
2104 0, /*tp_as_mapping*/
2105 0, /*tp_hash*/
2106 0, /*tp_call*/
2107 0, /*tp_str*/
2108 PyObject_GenericGetAttr,/*tp_getattro*/
2109 PyObject_GenericSetAttr,/*tp_setattro*/
2110 0, /*tp_as_buffer*/
2111 Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
2112 BZ2Decomp__doc__, /*tp_doc*/
2113 0, /*tp_traverse*/
2114 0, /*tp_clear*/
2115 0, /*tp_richcompare*/
2116 0, /*tp_weaklistoffset*/
2117 0, /*tp_iter*/
2118 0, /*tp_iternext*/
2119 BZ2Decomp_methods, /*tp_methods*/
2120 BZ2Decomp_members, /*tp_members*/
2121 0, /*tp_getset*/
2122 0, /*tp_base*/
2123 0, /*tp_dict*/
2124 0, /*tp_descr_get*/
2125 0, /*tp_descr_set*/
2126 0, /*tp_dictoffset*/
2127 (initproc)BZ2Decomp_init, /*tp_init*/
2128 PyType_GenericAlloc, /*tp_alloc*/
2129 PyType_GenericNew, /*tp_new*/
2130 _PyObject_Del, /*tp_free*/
2131 0, /*tp_is_gc*/
2132 };
2133
2134
2135 /* ===================================================================== */
2136 /* Module functions. */
2137
2138 PyDoc_STRVAR(bz2_compress__doc__,
2139 "compress(data [, compresslevel=9]) -> string\n\
2140 \n\
2141 Compress data in one shot. If you want to compress data sequentially,\n\
2142 use an instance of BZ2Compressor instead. The compresslevel parameter, if\n\
2143 given, must be a number between 1 and 9.\n\
2144 ");
2145
2146 static PyObject *
bz2_compress(PyObject * self,PyObject * args,PyObject * kwargs)2147 bz2_compress(PyObject *self, PyObject *args, PyObject *kwargs)
2148 {
2149 int compresslevel=9;
2150 int action;
2151 Py_buffer pdata;
2152 size_t input_left;
2153 size_t output_size = 0;
2154 PyObject *ret = NULL;
2155 bz_stream _bzs;
2156 bz_stream *bzs = &_bzs;
2157 int bzerror;
2158 static char *kwlist[] = {"data", "compresslevel", 0};
2159
2160 if (!PyArg_ParseTupleAndKeywords(args, kwargs, "s*|i",
2161 kwlist, &pdata,
2162 &compresslevel))
2163 return NULL;
2164
2165 if (compresslevel < 1 || compresslevel > 9) {
2166 PyErr_SetString(PyExc_ValueError,
2167 "compresslevel must be between 1 and 9");
2168 PyBuffer_Release(&pdata);
2169 return NULL;
2170 }
2171
2172 ret = PyString_FromStringAndSize(NULL, SMALLCHUNK);
2173 if (!ret) {
2174 PyBuffer_Release(&pdata);
2175 return NULL;
2176 }
2177
2178 memset(bzs, 0, sizeof(bz_stream));
2179
2180 bzs->next_in = pdata.buf;
2181 bzs->avail_in = MIN(pdata.len, UINT_MAX);
2182 input_left = pdata.len - bzs->avail_in;
2183
2184 bzs->next_out = BUF(ret);
2185 bzs->avail_out = PyString_GET_SIZE(ret);
2186
2187 bzerror = BZ2_bzCompressInit(bzs, compresslevel, 0, 0);
2188 if (bzerror != BZ_OK) {
2189 Util_CatchBZ2Error(bzerror);
2190 PyBuffer_Release(&pdata);
2191 Py_DECREF(ret);
2192 return NULL;
2193 }
2194
2195 action = input_left > 0 ? BZ_RUN : BZ_FINISH;
2196
2197 for (;;) {
2198 char *saved_next_out;
2199
2200 Py_BEGIN_ALLOW_THREADS
2201 saved_next_out = bzs->next_out;
2202 bzerror = BZ2_bzCompress(bzs, action);
2203 output_size += bzs->next_out - saved_next_out;
2204 Py_END_ALLOW_THREADS
2205
2206 if (bzerror == BZ_STREAM_END) {
2207 break;
2208 } else if (bzerror != BZ_RUN_OK && bzerror != BZ_FINISH_OK) {
2209 BZ2_bzCompressEnd(bzs);
2210 Util_CatchBZ2Error(bzerror);
2211 PyBuffer_Release(&pdata);
2212 Py_DECREF(ret);
2213 return NULL;
2214 }
2215 if (action == BZ_RUN && bzs->avail_in == 0) {
2216 if (input_left == 0) {
2217 action = BZ_FINISH;
2218 } else {
2219 bzs->avail_in = MIN(input_left, UINT_MAX);
2220 input_left -= bzs->avail_in;
2221 }
2222 }
2223 if (bzs->avail_out == 0) {
2224 size_t buffer_left = PyString_GET_SIZE(ret) - output_size;
2225 if (buffer_left == 0) {
2226 if (Util_GrowBuffer(&ret) < 0) {
2227 BZ2_bzCompressEnd(bzs);
2228 PyBuffer_Release(&pdata);
2229 return NULL;
2230 }
2231 bzs->next_out = BUF(ret) + output_size;
2232 buffer_left = PyString_GET_SIZE(ret) - output_size;
2233 }
2234 bzs->avail_out = MIN(buffer_left, UINT_MAX);
2235 }
2236 }
2237
2238 if (output_size != PyString_GET_SIZE(ret))
2239 _PyString_Resize(&ret, output_size); /* Sets ret to NULL on failure. */
2240
2241 BZ2_bzCompressEnd(bzs);
2242 PyBuffer_Release(&pdata);
2243 return ret;
2244 }
2245
2246 PyDoc_STRVAR(bz2_decompress__doc__,
2247 "decompress(data) -> decompressed data\n\
2248 \n\
2249 Decompress data in one shot. If you want to decompress data sequentially,\n\
2250 use an instance of BZ2Decompressor instead.\n\
2251 ");
2252
2253 static PyObject *
bz2_decompress(PyObject * self,PyObject * args)2254 bz2_decompress(PyObject *self, PyObject *args)
2255 {
2256 Py_buffer pdata;
2257 size_t input_left;
2258 size_t output_size = 0;
2259 PyObject *ret;
2260 bz_stream _bzs;
2261 bz_stream *bzs = &_bzs;
2262 int bzerror;
2263
2264 if (!PyArg_ParseTuple(args, "s*:decompress", &pdata))
2265 return NULL;
2266
2267 if (pdata.len == 0) {
2268 PyBuffer_Release(&pdata);
2269 return PyString_FromString("");
2270 }
2271
2272 ret = PyString_FromStringAndSize(NULL, SMALLCHUNK);
2273 if (!ret) {
2274 PyBuffer_Release(&pdata);
2275 return NULL;
2276 }
2277
2278 memset(bzs, 0, sizeof(bz_stream));
2279
2280 bzs->next_in = pdata.buf;
2281 bzs->avail_in = MIN(pdata.len, UINT_MAX);
2282 input_left = pdata.len - bzs->avail_in;
2283
2284 bzs->next_out = BUF(ret);
2285 bzs->avail_out = PyString_GET_SIZE(ret);
2286
2287 bzerror = BZ2_bzDecompressInit(bzs, 0, 0);
2288 if (bzerror != BZ_OK) {
2289 Util_CatchBZ2Error(bzerror);
2290 Py_DECREF(ret);
2291 PyBuffer_Release(&pdata);
2292 return NULL;
2293 }
2294
2295 for (;;) {
2296 char *saved_next_out;
2297
2298 Py_BEGIN_ALLOW_THREADS
2299 saved_next_out = bzs->next_out;
2300 bzerror = BZ2_bzDecompress(bzs);
2301 output_size += bzs->next_out - saved_next_out;
2302 Py_END_ALLOW_THREADS
2303
2304 if (bzerror == BZ_STREAM_END) {
2305 break;
2306 } else if (bzerror != BZ_OK) {
2307 BZ2_bzDecompressEnd(bzs);
2308 Util_CatchBZ2Error(bzerror);
2309 PyBuffer_Release(&pdata);
2310 Py_DECREF(ret);
2311 return NULL;
2312 }
2313 if (bzs->avail_in == 0) {
2314 if (input_left == 0) {
2315 BZ2_bzDecompressEnd(bzs);
2316 PyErr_SetString(PyExc_ValueError,
2317 "couldn't find end of stream");
2318 PyBuffer_Release(&pdata);
2319 Py_DECREF(ret);
2320 return NULL;
2321 }
2322 bzs->avail_in = MIN(input_left, UINT_MAX);
2323 input_left -= bzs->avail_in;
2324 }
2325 if (bzs->avail_out == 0) {
2326 size_t buffer_left = PyString_GET_SIZE(ret) - output_size;
2327 if (buffer_left == 0) {
2328 if (Util_GrowBuffer(&ret) < 0) {
2329 BZ2_bzDecompressEnd(bzs);
2330 PyBuffer_Release(&pdata);
2331 return NULL;
2332 }
2333 bzs->next_out = BUF(ret) + output_size;
2334 buffer_left = PyString_GET_SIZE(ret) - output_size;
2335 }
2336 bzs->avail_out = MIN(buffer_left, UINT_MAX);
2337 }
2338 }
2339
2340 if (output_size != PyString_GET_SIZE(ret))
2341 _PyString_Resize(&ret, output_size); /* Sets ret to NULL on failure. */
2342
2343 BZ2_bzDecompressEnd(bzs);
2344 PyBuffer_Release(&pdata);
2345 return ret;
2346 }
2347
2348 static PyMethodDef bz2_methods[] = {
2349 {"compress", (PyCFunction) bz2_compress, METH_VARARGS|METH_KEYWORDS,
2350 bz2_compress__doc__},
2351 {"decompress", (PyCFunction) bz2_decompress, METH_VARARGS,
2352 bz2_decompress__doc__},
2353 {NULL, NULL} /* sentinel */
2354 };
2355
2356 /* ===================================================================== */
2357 /* Initialization function. */
2358
2359 PyDoc_STRVAR(bz2__doc__,
2360 "The python bz2 module provides a comprehensive interface for\n\
2361 the bz2 compression library. It implements a complete file\n\
2362 interface, one shot (de)compression functions, and types for\n\
2363 sequential (de)compression.\n\
2364 ");
2365
2366 PyMODINIT_FUNC
initbz2(void)2367 initbz2(void)
2368 {
2369 PyObject *m;
2370
2371 if (PyType_Ready(&BZ2File_Type) < 0)
2372 return;
2373 if (PyType_Ready(&BZ2Comp_Type) < 0)
2374 return;
2375 if (PyType_Ready(&BZ2Decomp_Type) < 0)
2376 return;
2377
2378 m = Py_InitModule3("bz2", bz2_methods, bz2__doc__);
2379 if (m == NULL)
2380 return;
2381
2382 PyModule_AddObject(m, "__author__", PyString_FromString(__author__));
2383
2384 Py_INCREF(&BZ2File_Type);
2385 PyModule_AddObject(m, "BZ2File", (PyObject *)&BZ2File_Type);
2386
2387 Py_INCREF(&BZ2Comp_Type);
2388 PyModule_AddObject(m, "BZ2Compressor", (PyObject *)&BZ2Comp_Type);
2389
2390 Py_INCREF(&BZ2Decomp_Type);
2391 PyModule_AddObject(m, "BZ2Decompressor", (PyObject *)&BZ2Decomp_Type);
2392 }
2393