diff options
Diffstat (limited to 'Lib/bz2.py')
| -rw-r--r-- | Lib/bz2.py | 105 |
1 files changed, 70 insertions, 35 deletions
diff --git a/Lib/bz2.py b/Lib/bz2.py index a50adf79f7..c3075073f0 100644 --- a/Lib/bz2.py +++ b/Lib/bz2.py @@ -79,7 +79,8 @@ class BZ2File(io.BufferedIOBase): mode = "rb" mode_code = _MODE_READ self._decompressor = BZ2Decompressor() - self._buffer = None + self._buffer = b"" + self._buffer_offset = 0 elif mode in ("w", "wb"): mode = "wb" mode_code = _MODE_WRITE @@ -124,7 +125,8 @@ class BZ2File(io.BufferedIOBase): self._fp = None self._closefp = False self._mode = _MODE_CLOSED - self._buffer = None + self._buffer = b"" + self._buffer_offset = 0 @property def closed(self): @@ -157,15 +159,18 @@ class BZ2File(io.BufferedIOBase): raise ValueError("I/O operation on closed file") def _check_can_read(self): - if not self.readable(): + if self._mode not in (_MODE_READ, _MODE_READ_EOF): + self._check_not_closed() raise io.UnsupportedOperation("File not open for reading") def _check_can_write(self): - if not self.writable(): + if self._mode != _MODE_WRITE: + self._check_not_closed() raise io.UnsupportedOperation("File not open for writing") def _check_can_seek(self): - if not self.readable(): + if self._mode not in (_MODE_READ, _MODE_READ_EOF): + self._check_not_closed() raise io.UnsupportedOperation("Seeking is only supported " "on files open for reading") if not self._fp.seekable(): @@ -174,16 +179,13 @@ class BZ2File(io.BufferedIOBase): # Fill the readahead buffer if it is empty. Returns False on EOF. def _fill_buffer(self): + if self._mode == _MODE_READ_EOF: + return False # Depending on the input data, our call to the decompressor may not # return any data. In this case, try again after reading another block. - while True: - if self._buffer: - return True - - if self._decompressor.unused_data: - rawblock = self._decompressor.unused_data - else: - rawblock = self._fp.read(_BUFFER_SIZE) + while self._buffer_offset == len(self._buffer): + rawblock = (self._decompressor.unused_data or + self._fp.read(_BUFFER_SIZE)) if not rawblock: if self._decompressor.eof: @@ -199,30 +201,48 @@ class BZ2File(io.BufferedIOBase): self._decompressor = BZ2Decompressor() self._buffer = self._decompressor.decompress(rawblock) + self._buffer_offset = 0 + return True # Read data until EOF. # If return_data is false, consume the data without returning it. def _read_all(self, return_data=True): + # The loop assumes that _buffer_offset is 0. Ensure that this is true. + self._buffer = self._buffer[self._buffer_offset:] + self._buffer_offset = 0 + blocks = [] while self._fill_buffer(): if return_data: blocks.append(self._buffer) self._pos += len(self._buffer) - self._buffer = None + self._buffer = b"" if return_data: return b"".join(blocks) # Read a block of up to n bytes. # If return_data is false, consume the data without returning it. def _read_block(self, n, return_data=True): + # If we have enough data buffered, return immediately. + end = self._buffer_offset + n + if end <= len(self._buffer): + data = self._buffer[self._buffer_offset : end] + self._buffer_offset = end + self._pos += len(data) + return data if return_data else None + + # The loop assumes that _buffer_offset is 0. Ensure that this is true. + self._buffer = self._buffer[self._buffer_offset:] + self._buffer_offset = 0 + blocks = [] while n > 0 and self._fill_buffer(): if n < len(self._buffer): data = self._buffer[:n] - self._buffer = self._buffer[n:] + self._buffer_offset = n else: data = self._buffer - self._buffer = None + self._buffer = b"" if return_data: blocks.append(data) self._pos += len(data) @@ -238,9 +258,9 @@ class BZ2File(io.BufferedIOBase): """ with self._lock: self._check_can_read() - if self._mode == _MODE_READ_EOF or not self._fill_buffer(): + if not self._fill_buffer(): return b"" - return self._buffer + return self._buffer[self._buffer_offset:] def read(self, size=-1): """Read up to size uncompressed bytes from the file. @@ -250,7 +270,7 @@ class BZ2File(io.BufferedIOBase): """ with self._lock: self._check_can_read() - if self._mode == _MODE_READ_EOF or size == 0: + if size == 0: return b"" elif size < 0: return self._read_all() @@ -268,15 +288,19 @@ class BZ2File(io.BufferedIOBase): # In this case we make multiple reads, to avoid returning b"". with self._lock: self._check_can_read() - if (size == 0 or self._mode == _MODE_READ_EOF or - not self._fill_buffer()): + if (size == 0 or + # Only call _fill_buffer() if the buffer is actually empty. + # This gives a significant speedup if *size* is small. + (self._buffer_offset == len(self._buffer) and not self._fill_buffer())): return b"" - if 0 < size < len(self._buffer): - data = self._buffer[:size] - self._buffer = self._buffer[size:] + if size > 0: + data = self._buffer[self._buffer_offset : + self._buffer_offset + size] + self._buffer_offset += len(data) else: - data = self._buffer - self._buffer = None + data = self._buffer[self._buffer_offset:] + self._buffer = b"" + self._buffer_offset = 0 self._pos += len(data) return data @@ -295,10 +319,20 @@ class BZ2File(io.BufferedIOBase): non-negative, no more than size bytes will be read (in which case the line may be incomplete). Returns b'' if already at EOF. """ - if not hasattr(size, "__index__"): - raise TypeError("Integer argument expected") - size = size.__index__() + if not isinstance(size, int): + if not hasattr(size, "__index__"): + raise TypeError("Integer argument expected") + size = size.__index__() with self._lock: + self._check_can_read() + # Shortcut for the common case - the whole line is in the buffer. + if size < 0: + end = self._buffer.find(b"\n", self._buffer_offset) + 1 + if end > 0: + line = self._buffer[self._buffer_offset : end] + self._buffer_offset = end + self._pos += len(line) + return line return io.BufferedIOBase.readline(self, size) def readlines(self, size=-1): @@ -308,9 +342,10 @@ class BZ2File(io.BufferedIOBase): further lines will be read once the total size of the lines read so far equals or exceeds size. """ - if not hasattr(size, "__index__"): - raise TypeError("Integer argument expected") - size = size.__index__() + if not isinstance(size, int): + if not hasattr(size, "__index__"): + raise TypeError("Integer argument expected") + size = size.__index__() with self._lock: return io.BufferedIOBase.readlines(self, size) @@ -345,7 +380,8 @@ class BZ2File(io.BufferedIOBase): self._mode = _MODE_READ self._pos = 0 self._decompressor = BZ2Decompressor() - self._buffer = None + self._buffer = b"" + self._buffer_offset = 0 def seek(self, offset, whence=0): """Change the file position. @@ -385,8 +421,7 @@ class BZ2File(io.BufferedIOBase): offset -= self._pos # Read and discard data until we reach the desired position. - if self._mode != _MODE_READ_EOF: - self._read_block(offset, return_data=False) + self._read_block(offset, return_data=False) return self._pos |