Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 1 | """Interface to the libbzip2 compression library. |
| 2 | |
| 3 | This module provides a file interface, classes for incremental |
| 4 | (de)compression, and functions for one-shot (de)compression. |
| 5 | """ |
| 6 | |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 7 | __all__ = ["BZ2File", "BZ2Compressor", "BZ2Decompressor", |
| 8 | "open", "compress", "decompress"] |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 9 | |
| 10 | __author__ = "Nadeem Vawda <nadeem.vawda@gmail.com>" |
| 11 | |
| 12 | import io |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 13 | import warnings |
| 14 | |
Nadeem Vawda | 72750a8 | 2012-01-18 01:57:14 +0200 | [diff] [blame] | 15 | try: |
| 16 | from threading import RLock |
| 17 | except ImportError: |
| 18 | from dummy_threading import RLock |
| 19 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 20 | from _bz2 import BZ2Compressor, BZ2Decompressor |
| 21 | |
| 22 | |
| 23 | _MODE_CLOSED = 0 |
| 24 | _MODE_READ = 1 |
| 25 | _MODE_READ_EOF = 2 |
| 26 | _MODE_WRITE = 3 |
| 27 | |
| 28 | _BUFFER_SIZE = 8192 |
| 29 | |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 30 | _builtin_open = open |
| 31 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 32 | |
| 33 | class BZ2File(io.BufferedIOBase): |
| 34 | |
| 35 | """A file object providing transparent bzip2 (de)compression. |
| 36 | |
| 37 | A BZ2File can act as a wrapper for an existing file object, or refer |
| 38 | directly to a named file on disk. |
| 39 | |
| 40 | Note that BZ2File provides a *binary* file interface - data read is |
| 41 | returned as bytes, and data to be written should be given as bytes. |
| 42 | """ |
| 43 | |
Nadeem Vawda | aebcdba | 2012-06-04 23:31:20 +0200 | [diff] [blame] | 44 | def __init__(self, filename, mode="r", buffering=None, compresslevel=9): |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 45 | """Open a bzip2-compressed file. |
| 46 | |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 47 | If filename is a str or bytes object, it gives the name |
| 48 | of the file to be opened. Otherwise, it should be a file object, |
| 49 | which will be used to read or write the compressed data. |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 50 | |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 51 | mode can be 'r' for reading (default), 'w' for (over)writing, |
| 52 | or 'a' for appending. These can equivalently be given as 'rb', |
| 53 | 'wb', and 'ab'. |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 54 | |
| 55 | buffering is ignored. Its use is deprecated. |
| 56 | |
Nadeem Vawda | cac8909 | 2012-02-04 13:08:11 +0200 | [diff] [blame] | 57 | If mode is 'w' or 'a', compresslevel can be a number between 1 |
| 58 | and 9 specifying the level of compression: 1 produces the least |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 59 | compression, and 9 (default) produces the most compression. |
Nadeem Vawda | cac8909 | 2012-02-04 13:08:11 +0200 | [diff] [blame] | 60 | |
| 61 | If mode is 'r', the input file may be the concatenation of |
| 62 | multiple compressed streams. |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 63 | """ |
| 64 | # This lock must be recursive, so that BufferedIOBase's |
| 65 | # readline(), readlines() and writelines() don't deadlock. |
Nadeem Vawda | 72750a8 | 2012-01-18 01:57:14 +0200 | [diff] [blame] | 66 | self._lock = RLock() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 67 | self._fp = None |
| 68 | self._closefp = False |
| 69 | self._mode = _MODE_CLOSED |
| 70 | self._pos = 0 |
| 71 | self._size = -1 |
| 72 | |
| 73 | if buffering is not None: |
| 74 | warnings.warn("Use of 'buffering' argument is deprecated", |
| 75 | DeprecationWarning) |
| 76 | |
| 77 | if not (1 <= compresslevel <= 9): |
| 78 | raise ValueError("compresslevel must be between 1 and 9") |
| 79 | |
| 80 | if mode in ("", "r", "rb"): |
| 81 | mode = "rb" |
| 82 | mode_code = _MODE_READ |
| 83 | self._decompressor = BZ2Decompressor() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 84 | self._buffer = b"" |
| 85 | self._buffer_offset = 0 |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 86 | elif mode in ("w", "wb"): |
| 87 | mode = "wb" |
| 88 | mode_code = _MODE_WRITE |
Nadeem Vawda | 249ab5e | 2011-09-11 22:38:11 +0200 | [diff] [blame] | 89 | self._compressor = BZ2Compressor(compresslevel) |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 90 | elif mode in ("a", "ab"): |
| 91 | mode = "ab" |
| 92 | mode_code = _MODE_WRITE |
Nadeem Vawda | 249ab5e | 2011-09-11 22:38:11 +0200 | [diff] [blame] | 93 | self._compressor = BZ2Compressor(compresslevel) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 94 | else: |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 95 | raise ValueError("Invalid mode: %r" % (mode,)) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 96 | |
Nadeem Vawda | aebcdba | 2012-06-04 23:31:20 +0200 | [diff] [blame] | 97 | if isinstance(filename, (str, bytes)): |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 98 | self._fp = _builtin_open(filename, mode) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 99 | self._closefp = True |
| 100 | self._mode = mode_code |
Nadeem Vawda | aebcdba | 2012-06-04 23:31:20 +0200 | [diff] [blame] | 101 | elif hasattr(filename, "read") or hasattr(filename, "write"): |
| 102 | self._fp = filename |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 103 | self._mode = mode_code |
| 104 | else: |
Nadeem Vawda | aebcdba | 2012-06-04 23:31:20 +0200 | [diff] [blame] | 105 | raise TypeError("filename must be a str or bytes object, or a file") |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 106 | |
| 107 | def close(self): |
| 108 | """Flush and close the file. |
| 109 | |
| 110 | May be called more than once without error. Once the file is |
| 111 | closed, any other operation on it will raise a ValueError. |
| 112 | """ |
| 113 | with self._lock: |
| 114 | if self._mode == _MODE_CLOSED: |
| 115 | return |
| 116 | try: |
| 117 | if self._mode in (_MODE_READ, _MODE_READ_EOF): |
| 118 | self._decompressor = None |
| 119 | elif self._mode == _MODE_WRITE: |
| 120 | self._fp.write(self._compressor.flush()) |
| 121 | self._compressor = None |
| 122 | finally: |
Antoine Pitrou | 24ce386 | 2011-04-03 17:08:49 +0200 | [diff] [blame] | 123 | try: |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 124 | if self._closefp: |
| 125 | self._fp.close() |
| 126 | finally: |
| 127 | self._fp = None |
| 128 | self._closefp = False |
| 129 | self._mode = _MODE_CLOSED |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 130 | self._buffer = b"" |
| 131 | self._buffer_offset = 0 |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 132 | |
| 133 | @property |
| 134 | def closed(self): |
| 135 | """True if this file is closed.""" |
| 136 | return self._mode == _MODE_CLOSED |
| 137 | |
| 138 | def fileno(self): |
| 139 | """Return the file descriptor for the underlying file.""" |
Nadeem Vawda | 44ae4a2 | 2011-11-30 17:39:30 +0200 | [diff] [blame] | 140 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 141 | return self._fp.fileno() |
| 142 | |
| 143 | def seekable(self): |
| 144 | """Return whether the file supports seeking.""" |
Nadeem Vawda | ae557d7 | 2012-02-12 01:51:38 +0200 | [diff] [blame] | 145 | return self.readable() and self._fp.seekable() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 146 | |
| 147 | def readable(self): |
| 148 | """Return whether the file was opened for reading.""" |
Nadeem Vawda | 44ae4a2 | 2011-11-30 17:39:30 +0200 | [diff] [blame] | 149 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 150 | return self._mode in (_MODE_READ, _MODE_READ_EOF) |
| 151 | |
| 152 | def writable(self): |
| 153 | """Return whether the file was opened for writing.""" |
Nadeem Vawda | 44ae4a2 | 2011-11-30 17:39:30 +0200 | [diff] [blame] | 154 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 155 | return self._mode == _MODE_WRITE |
| 156 | |
| 157 | # Mode-checking helper functions. |
| 158 | |
| 159 | def _check_not_closed(self): |
| 160 | if self.closed: |
| 161 | raise ValueError("I/O operation on closed file") |
| 162 | |
| 163 | def _check_can_read(self): |
Nadeem Vawda | b7a0bfe | 2012-09-30 23:58:01 +0200 | [diff] [blame] | 164 | if self._mode not in (_MODE_READ, _MODE_READ_EOF): |
Nadeem Vawda | 452add0 | 2012-10-01 23:02:50 +0200 | [diff] [blame] | 165 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 166 | raise io.UnsupportedOperation("File not open for reading") |
| 167 | |
| 168 | def _check_can_write(self): |
Nadeem Vawda | b7a0bfe | 2012-09-30 23:58:01 +0200 | [diff] [blame] | 169 | if self._mode != _MODE_WRITE: |
Nadeem Vawda | 452add0 | 2012-10-01 23:02:50 +0200 | [diff] [blame] | 170 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 171 | raise io.UnsupportedOperation("File not open for writing") |
| 172 | |
| 173 | def _check_can_seek(self): |
Nadeem Vawda | b7a0bfe | 2012-09-30 23:58:01 +0200 | [diff] [blame] | 174 | if self._mode not in (_MODE_READ, _MODE_READ_EOF): |
Nadeem Vawda | 452add0 | 2012-10-01 23:02:50 +0200 | [diff] [blame] | 175 | self._check_not_closed() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 176 | raise io.UnsupportedOperation("Seeking is only supported " |
Nadeem Vawda | f1a1af2 | 2011-05-25 00:32:08 +0200 | [diff] [blame] | 177 | "on files open for reading") |
Nadeem Vawda | ae557d7 | 2012-02-12 01:51:38 +0200 | [diff] [blame] | 178 | if not self._fp.seekable(): |
| 179 | raise io.UnsupportedOperation("The underlying file object " |
| 180 | "does not support seeking") |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 181 | |
| 182 | # Fill the readahead buffer if it is empty. Returns False on EOF. |
| 183 | def _fill_buffer(self): |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 184 | if self._mode == _MODE_READ_EOF: |
| 185 | return False |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 186 | # Depending on the input data, our call to the decompressor may not |
| 187 | # return any data. In this case, try again after reading another block. |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 188 | while self._buffer_offset == len(self._buffer): |
| 189 | rawblock = (self._decompressor.unused_data or |
| 190 | self._fp.read(_BUFFER_SIZE)) |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 191 | |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 192 | if not rawblock: |
| 193 | if self._decompressor.eof: |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 194 | # End-of-stream marker and end of file. We're good. |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 195 | self._mode = _MODE_READ_EOF |
| 196 | self._size = self._pos |
| 197 | return False |
| 198 | else: |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 199 | # Problem - we were expecting more compressed data. |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 200 | raise EOFError("Compressed file ended before the " |
| 201 | "end-of-stream marker was reached") |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 202 | |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 203 | if self._decompressor.eof: |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 204 | # Continue to next stream. |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 205 | self._decompressor = BZ2Decompressor() |
| 206 | |
| 207 | self._buffer = self._decompressor.decompress(rawblock) |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 208 | self._buffer_offset = 0 |
| 209 | return True |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 210 | |
| 211 | # Read data until EOF. |
| 212 | # If return_data is false, consume the data without returning it. |
| 213 | def _read_all(self, return_data=True): |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 214 | # The loop assumes that _buffer_offset is 0. Ensure that this is true. |
| 215 | self._buffer = self._buffer[self._buffer_offset:] |
| 216 | self._buffer_offset = 0 |
| 217 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 218 | blocks = [] |
| 219 | while self._fill_buffer(): |
| 220 | if return_data: |
| 221 | blocks.append(self._buffer) |
| 222 | self._pos += len(self._buffer) |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 223 | self._buffer = b"" |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 224 | if return_data: |
| 225 | return b"".join(blocks) |
| 226 | |
| 227 | # Read a block of up to n bytes. |
| 228 | # If return_data is false, consume the data without returning it. |
| 229 | def _read_block(self, n, return_data=True): |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 230 | # If we have enough data buffered, return immediately. |
| 231 | end = self._buffer_offset + n |
| 232 | if end <= len(self._buffer): |
| 233 | data = self._buffer[self._buffer_offset : end] |
| 234 | self._buffer_offset = end |
| 235 | self._pos += len(data) |
Nadeem Vawda | 9e2a28e | 2012-09-30 13:41:29 +0200 | [diff] [blame] | 236 | return data if return_data else None |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 237 | |
| 238 | # The loop assumes that _buffer_offset is 0. Ensure that this is true. |
| 239 | self._buffer = self._buffer[self._buffer_offset:] |
| 240 | self._buffer_offset = 0 |
| 241 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 242 | blocks = [] |
| 243 | while n > 0 and self._fill_buffer(): |
| 244 | if n < len(self._buffer): |
| 245 | data = self._buffer[:n] |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 246 | self._buffer_offset = n |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 247 | else: |
| 248 | data = self._buffer |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 249 | self._buffer = b"" |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 250 | if return_data: |
| 251 | blocks.append(data) |
| 252 | self._pos += len(data) |
| 253 | n -= len(data) |
| 254 | if return_data: |
| 255 | return b"".join(blocks) |
| 256 | |
| 257 | def peek(self, n=0): |
| 258 | """Return buffered data without advancing the file position. |
| 259 | |
| 260 | Always returns at least one byte of data, unless at EOF. |
| 261 | The exact number of bytes returned is unspecified. |
| 262 | """ |
| 263 | with self._lock: |
| 264 | self._check_can_read() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 265 | if not self._fill_buffer(): |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 266 | return b"" |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 267 | return self._buffer[self._buffer_offset:] |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 268 | |
| 269 | def read(self, size=-1): |
| 270 | """Read up to size uncompressed bytes from the file. |
| 271 | |
| 272 | If size is negative or omitted, read until EOF is reached. |
| 273 | Returns b'' if the file is already at EOF. |
| 274 | """ |
| 275 | with self._lock: |
| 276 | self._check_can_read() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 277 | if size == 0: |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 278 | return b"" |
| 279 | elif size < 0: |
| 280 | return self._read_all() |
| 281 | else: |
| 282 | return self._read_block(size) |
| 283 | |
| 284 | def read1(self, size=-1): |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 285 | """Read up to size uncompressed bytes, while trying to avoid |
| 286 | making multiple reads from the underlying stream. |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 287 | |
| 288 | Returns b'' if the file is at EOF. |
| 289 | """ |
Nadeem Vawda | 8280b4b | 2012-08-04 15:29:28 +0200 | [diff] [blame] | 290 | # Usually, read1() calls _fp.read() at most once. However, sometimes |
| 291 | # this does not give enough data for the decompressor to make progress. |
| 292 | # In this case we make multiple reads, to avoid returning b"". |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 293 | with self._lock: |
| 294 | self._check_can_read() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 295 | if (size == 0 or |
| 296 | # Only call _fill_buffer() if the buffer is actually empty. |
| 297 | # This gives a significant speedup if *size* is small. |
| 298 | (self._buffer_offset == len(self._buffer) and not self._fill_buffer())): |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 299 | return b"" |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 300 | if size > 0: |
| 301 | data = self._buffer[self._buffer_offset : |
| 302 | self._buffer_offset + size] |
| 303 | self._buffer_offset += len(data) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 304 | else: |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 305 | data = self._buffer[self._buffer_offset:] |
| 306 | self._buffer = b"" |
| 307 | self._buffer_offset = 0 |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 308 | self._pos += len(data) |
| 309 | return data |
| 310 | |
| 311 | def readinto(self, b): |
| 312 | """Read up to len(b) bytes into b. |
Antoine Pitrou | 24ce386 | 2011-04-03 17:08:49 +0200 | [diff] [blame] | 313 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 314 | Returns the number of bytes read (0 for EOF). |
| 315 | """ |
| 316 | with self._lock: |
| 317 | return io.BufferedIOBase.readinto(self, b) |
| 318 | |
| 319 | def readline(self, size=-1): |
| 320 | """Read a line of uncompressed bytes from the file. |
| 321 | |
| 322 | The terminating newline (if present) is retained. If size is |
| 323 | non-negative, no more than size bytes will be read (in which |
| 324 | case the line may be incomplete). Returns b'' if already at EOF. |
| 325 | """ |
Nadeem Vawda | eb70be2 | 2012-10-01 23:05:32 +0200 | [diff] [blame] | 326 | if not isinstance(size, int): |
| 327 | if not hasattr(size, "__index__"): |
| 328 | raise TypeError("Integer argument expected") |
| 329 | size = size.__index__() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 330 | with self._lock: |
Nadeem Vawda | 138ad50 | 2012-10-01 23:04:11 +0200 | [diff] [blame] | 331 | self._check_can_read() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 332 | # Shortcut for the common case - the whole line is in the buffer. |
| 333 | if size < 0: |
| 334 | end = self._buffer.find(b"\n", self._buffer_offset) + 1 |
| 335 | if end > 0: |
| 336 | line = self._buffer[self._buffer_offset : end] |
| 337 | self._buffer_offset = end |
| 338 | self._pos += len(line) |
| 339 | return line |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 340 | return io.BufferedIOBase.readline(self, size) |
| 341 | |
| 342 | def readlines(self, size=-1): |
| 343 | """Read a list of lines of uncompressed bytes from the file. |
| 344 | |
| 345 | size can be specified to control the number of lines read: no |
| 346 | further lines will be read once the total size of the lines read |
| 347 | so far equals or exceeds size. |
| 348 | """ |
Nadeem Vawda | eb70be2 | 2012-10-01 23:05:32 +0200 | [diff] [blame] | 349 | if not isinstance(size, int): |
| 350 | if not hasattr(size, "__index__"): |
| 351 | raise TypeError("Integer argument expected") |
| 352 | size = size.__index__() |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 353 | with self._lock: |
| 354 | return io.BufferedIOBase.readlines(self, size) |
| 355 | |
| 356 | def write(self, data): |
| 357 | """Write a byte string to the file. |
| 358 | |
| 359 | Returns the number of uncompressed bytes written, which is |
| 360 | always len(data). Note that due to buffering, the file on disk |
| 361 | may not reflect the data written until close() is called. |
| 362 | """ |
| 363 | with self._lock: |
| 364 | self._check_can_write() |
| 365 | compressed = self._compressor.compress(data) |
| 366 | self._fp.write(compressed) |
| 367 | self._pos += len(data) |
| 368 | return len(data) |
| 369 | |
| 370 | def writelines(self, seq): |
| 371 | """Write a sequence of byte strings to the file. |
| 372 | |
| 373 | Returns the number of uncompressed bytes written. |
| 374 | seq can be any iterable yielding byte strings. |
| 375 | |
| 376 | Line separators are not added between the written byte strings. |
| 377 | """ |
| 378 | with self._lock: |
| 379 | return io.BufferedIOBase.writelines(self, seq) |
| 380 | |
| 381 | # Rewind the file to the beginning of the data stream. |
| 382 | def _rewind(self): |
| 383 | self._fp.seek(0, 0) |
| 384 | self._mode = _MODE_READ |
| 385 | self._pos = 0 |
| 386 | self._decompressor = BZ2Decompressor() |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 387 | self._buffer = b"" |
| 388 | self._buffer_offset = 0 |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 389 | |
| 390 | def seek(self, offset, whence=0): |
| 391 | """Change the file position. |
| 392 | |
| 393 | The new position is specified by offset, relative to the |
| 394 | position indicated by whence. Values for whence are: |
| 395 | |
| 396 | 0: start of stream (default); offset must not be negative |
| 397 | 1: current stream position |
| 398 | 2: end of stream; offset must not be positive |
| 399 | |
| 400 | Returns the new file position. |
| 401 | |
| 402 | Note that seeking is emulated, so depending on the parameters, |
| 403 | this operation may be extremely slow. |
| 404 | """ |
| 405 | with self._lock: |
| 406 | self._check_can_seek() |
| 407 | |
| 408 | # Recalculate offset as an absolute file position. |
| 409 | if whence == 0: |
| 410 | pass |
| 411 | elif whence == 1: |
| 412 | offset = self._pos + offset |
| 413 | elif whence == 2: |
| 414 | # Seeking relative to EOF - we need to know the file's size. |
| 415 | if self._size < 0: |
| 416 | self._read_all(return_data=False) |
| 417 | offset = self._size + offset |
| 418 | else: |
Nadeem Vawda | 3b4a4f5 | 2012-10-08 19:20:49 +0200 | [diff] [blame] | 419 | raise ValueError("Invalid value for whence: %s" % (whence,)) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 420 | |
| 421 | # Make it so that offset is the number of bytes to skip forward. |
| 422 | if offset < self._pos: |
| 423 | self._rewind() |
| 424 | else: |
| 425 | offset -= self._pos |
| 426 | |
| 427 | # Read and discard data until we reach the desired position. |
Nadeem Vawda | 6c57318 | 2012-09-30 03:57:33 +0200 | [diff] [blame] | 428 | self._read_block(offset, return_data=False) |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 429 | |
| 430 | return self._pos |
| 431 | |
| 432 | def tell(self): |
| 433 | """Return the current file position.""" |
| 434 | with self._lock: |
| 435 | self._check_not_closed() |
| 436 | return self._pos |
| 437 | |
| 438 | |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 439 | def open(filename, mode="rb", compresslevel=9, |
| 440 | encoding=None, errors=None, newline=None): |
| 441 | """Open a bzip2-compressed file in binary or text mode. |
| 442 | |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 443 | The filename argument can be an actual filename (a str or bytes |
| 444 | object), or an existing file object to read from or write to. |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 445 | |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 446 | The mode argument can be "r", "rb", "w", "wb", "a" or "ab" for |
| 447 | binary mode, or "rt", "wt" or "at" for text mode. The default mode |
| 448 | is "rb", and the default compresslevel is 9. |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 449 | |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 450 | For binary mode, this function is equivalent to the BZ2File |
| 451 | constructor: BZ2File(filename, mode, compresslevel). In this case, |
| 452 | the encoding, errors and newline arguments must not be provided. |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 453 | |
| 454 | For text mode, a BZ2File object is created, and wrapped in an |
Nadeem Vawda | 4907b0a | 2012-10-08 20:31:34 +0200 | [diff] [blame] | 455 | io.TextIOWrapper instance with the specified encoding, error |
| 456 | handling behavior, and line ending(s). |
Nadeem Vawda | af518c1 | 2012-06-04 23:32:38 +0200 | [diff] [blame] | 457 | |
| 458 | """ |
| 459 | if "t" in mode: |
| 460 | if "b" in mode: |
| 461 | raise ValueError("Invalid mode: %r" % (mode,)) |
| 462 | else: |
| 463 | if encoding is not None: |
| 464 | raise ValueError("Argument 'encoding' not supported in binary mode") |
| 465 | if errors is not None: |
| 466 | raise ValueError("Argument 'errors' not supported in binary mode") |
| 467 | if newline is not None: |
| 468 | raise ValueError("Argument 'newline' not supported in binary mode") |
| 469 | |
| 470 | bz_mode = mode.replace("t", "") |
| 471 | binary_file = BZ2File(filename, bz_mode, compresslevel=compresslevel) |
| 472 | |
| 473 | if "t" in mode: |
| 474 | return io.TextIOWrapper(binary_file, encoding, errors, newline) |
| 475 | else: |
| 476 | return binary_file |
| 477 | |
| 478 | |
Antoine Pitrou | 37dc5f8 | 2011-04-03 17:05:46 +0200 | [diff] [blame] | 479 | def compress(data, compresslevel=9): |
| 480 | """Compress a block of data. |
| 481 | |
| 482 | compresslevel, if given, must be a number between 1 and 9. |
| 483 | |
| 484 | For incremental compression, use a BZ2Compressor object instead. |
| 485 | """ |
| 486 | comp = BZ2Compressor(compresslevel) |
| 487 | return comp.compress(data) + comp.flush() |
| 488 | |
| 489 | |
| 490 | def decompress(data): |
| 491 | """Decompress a block of data. |
| 492 | |
| 493 | For incremental decompression, use a BZ2Decompressor object instead. |
| 494 | """ |
| 495 | if len(data) == 0: |
| 496 | return b"" |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 497 | |
Nadeem Vawda | 98838ba | 2011-05-30 01:12:24 +0200 | [diff] [blame] | 498 | results = [] |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 499 | while True: |
| 500 | decomp = BZ2Decompressor() |
Nadeem Vawda | 98838ba | 2011-05-30 01:12:24 +0200 | [diff] [blame] | 501 | results.append(decomp.decompress(data)) |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 502 | if not decomp.eof: |
| 503 | raise ValueError("Compressed data ended before the " |
| 504 | "end-of-stream marker was reached") |
| 505 | if not decomp.unused_data: |
Nadeem Vawda | 98838ba | 2011-05-30 01:12:24 +0200 | [diff] [blame] | 506 | return b"".join(results) |
Nadeem Vawda | 55b4338 | 2011-05-27 01:52:15 +0200 | [diff] [blame] | 507 | # There is unused data left over. Proceed to next stream. |
| 508 | data = decomp.unused_data |