[issue5863] bz2.BZ2File should accept other file-like objects.

Lib/bz2.py

+"""Interface to the libbzip2 compression library.
+
+This module provides a file interface, classes for incremental
+(de)compression, and functions for one-shot (de)compression.
+"""
+
+__all__ = ["BZ2File", "BZ2Compressor", "BZ2Decompressor", "compress",
+           "decompress"]
+
+__author__ = "Nadeem Vawda <nadeem.vawda@gmail.com>"
+
+import _bz2
+import threading
+import warnings
+
+
+_MODE_CLOSED   = 0
+_MODE_READ     = 1
+_MODE_READ_EOF = 2
+_MODE_WRITE    = 3
+
+_BUFFER_SIZE = 8192
+
+
+class BZ2File:

[Antoine Pitrou, 2011/03/13 17:31:28]

Is there any reason it doesn't inherit io.BufferedIOBase?
(it should also bring you a couple of methods implemented for free: readlines,
writelines, __iter__, __next__, __enter__, __exit__)

You should probably also implement fileno() (simply return self.fp.fileno()) and
the `closed` property.

+
+    """A file object providing transparent bzip2 (de)compression.
+
+    A BZ2File can act as a wrapper for an existing file object, or refer
+    directly to a named file on disk.
+
+    Note that BZ2File provides a *binary* file interface - data read is
+    returned as bytes, and data to be written should be given as bytes.
+    """
+
+    def __init__(self, filename=None, mode="r", buffering=None,
+                 compresslevel=9, fileobj=None):
+        """Open a bzip2-compressed file.
+
+        If filename is given, open the named file. Otherwise, operate on
+        the file object given by fileobj. Exactly one of these two
+        parameters should be provided.
+
+        mode can be 'r' for reading (default), or 'w' for writing.
+
+        buffering is ignored. Its use is deprecated.
+
+        If mode is 'w', compresslevel can be a number between 1 and 9
+        specifying the level of compression: 1 produces the least
+        compression, and 9 (default) produces the most compression.
+        """
+        self._lock = threading.Lock()
+        self._fp = None
+        self._closefp = False
+        self._mode = _MODE_CLOSED
+        self._pos = 0
+        self._size = -1
+
+        if buffering is not None:
+            warnings.warn("Use of 'buffering' argument is deprecated",
+                          DeprecationWarning)
+
+        if not (1 <= compresslevel <= 9):
+            raise ValueError("compresslevel must be between 1 and 9")
+
+        if mode in ("", "r", "rb"):
+            mode = "rb"
+            mode_code = _MODE_READ
+            self._decompressor = _bz2.Decompressor()
+            self._buffer = None
+        elif mode in ("w", "wb"):
+            mode = "wb"
+            mode_code = _MODE_WRITE
+            self._compressor = _bz2.Compressor()
+        else:
+            raise ValueError("Invalid mode: {!r}".format(mode))
+
+        if filename is not None and fileobj is None:
+            self._fp = open(filename, mode)
+            self._closefp = True
+            self._mode = mode_code
+        elif fileobj is not None and filename is None:
+            self._fp = fileobj
+            self._mode = mode_code
+        else:
+            raise ValueError("Must give exactly one of filename and fileobj")
+
+    def __del__(self):
+        self._close()
+
+    def _close(self):
+        if self._mode == _MODE_CLOSED:
+            return
+        try:
+            if self._mode in (_MODE_READ, _MODE_READ_EOF):
+                self._decompressor = None
+            elif self._mode == _MODE_WRITE:
+                self._fp.write(self._compressor.flush())
+                self._compressor = None
+        finally:
+            try: 
+                if self._closefp:
+                   self._fp.close()
+            finally:
+                self._fp = None
+                self._closefp = False
+                self._mode = _MODE_CLOSED
+                self._buffer = None
+
+    # Context management protocol:
+
+    def __enter__(self):
+        self._check_not_closed()
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+    # Iterator protocol for line iteration:
+
+    def __iter__(self):
+        self._check_not_closed()
+        return self
+
+    def __next__(self):
+        with self._lock:
+            self._check_can_read()
+            line = self._read_line(-1)
+            if not line:
+                raise StopIteration()
+            return line
+
+    # Public I/O methods:
+
+    def read(self, size=-1):
+        """Read up to size uncompressed bytes from the file.
+
+        If size is negative or omitted, read until EOF is reached.
+        Returns b'' if the file is already at EOF.
+        """
+        with self._lock:
+            self._check_can_read()
+            if self._mode == _MODE_READ_EOF or size == 0:
+                return b""
+            elif size < 0:
+                return self._read_all()
+            else:
+                return self._read_block(size)
+
+    def peek(self, n=0):
+        """Return buffered data without advancing the file position.
+
+        At least one byte of data will be returned (unless at EOF).
+        The exact number of bytes returned is unspecified.
+        """
+        with self._lock:
+            self._check_can_read()
+            if self._mode == _MODE_READ_EOF:
+                return b""
+            self._fill_buffer()
+            return self._buffer
+
+    def readline(self, size=-1):
+        """Read a line of uncompressed bytes from the file.
+
+        The terminating newline (if present) is retained. If size is
+        non-negative, no more than size bytes will be read (in which
+        case the line may be incomplete). Returns b'' if already at EOF.
+        """
+        with self._lock:
+            self._check_can_read()
+            return self._read_line(size)
+
+    def readlines(self, size=-1):
+        """Read a list of lines of uncompressed bytes from the file.
+
+        size can be specified to control the number of lines read: no
+        further lines will be read once the total size of the lines read
+        so far equals or exceeds size.
+        """
+        with self._lock:
+            self._check_can_read()
+            lines = []
+            nread = 0
+            while (size < 0 or nread < size) and self._mode != _MODE_READ_EOF:
+                line = self._read_line(-1)
+                if not line:
+                    break
+                lines.append(line)
+                nread += len(line)
+            return lines
+
+    def write(self, data):
+        """Write a byte string to the file.
+
+        Returns the number of uncompressed bytes written, which is
+        always len(data). Note that due to buffering, the file on disk
+        may not reflect the data written until close() is called.
+        """
+        with self._lock:
+            self._check_can_write()
+            return self._write(data)
+
+    def writelines(self, seq):
+        """Write a sequence of byte strings to the file.
+
+        Returns the number of uncompressed bytes written.
+        seq can be any iterable yielding byte strings.
+
+        Line separators are not added between the written byte strings.
+        """
+        with self._lock:
+            self._check_can_write()
+            nwritten = 0
+            for block in seq:
+                nwritten += self._write(block)
+            return nwritten
+
+    def seek(self, offset, whence=0):
+        """Change the file position.
+
+        The new position is specified by offset, relative to the
+        position indicated by whence. Values for whence are:
+
+            0: start of stream (default); offset must not be negative
+            1: current stream position
+            2: end of stream; offset must not be positive
+
+        Returns the new file position.
+
+        Note that seeking is emulated, so depending on the parameters,
+        this operation may be extremely slow.
+        """
+        with self._lock:
+            self._check_can_seek()
+
+            # Recalculate offset as an absolute file position.
+            if whence == 0:
+                pass
+            elif whence == 1:
+                offset = self._pos + offset
+            elif whence == 2:
+                # Seeking relative to EOF - we need to know the file's size.
+                if self._size < 0:
+                    self._read_all(return_data=False)
+                offset = self._size + offset
+            else:
+                raise ValueError("Invalid value for whence: {}".format(whence))
+
+            # Make it so that offset is the number of bytes to skip forward.
+            if offset < self._pos:
+                self._rewind()
+            else:
+                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)
+
+            return self._pos
+
+    def tell(self):
+        """Return the current file position."""
+        with self._lock:
+            self._check_not_closed()
+            return self._pos
+
+    def close(self):
+        """Flush and close the file.
+
+        May be called more than once without error. Once the file is
+        closed, any other operation on it will raise a ValueError.
+        """
+        with self._lock:
+            self._close()
+
+    @property
+    def closed(self):
+        """True if this file is closed."""
+        return self._mode == _MODE_CLOSED
+
+    # Helper methods. All methods below here (except for _check_*())
+    # assume that the caller owns the object's lock.
+
+    def _check_not_closed(self):
+        if self._mode == _MODE_CLOSED:
+            raise ValueError("I/O operation on closed file")
+
+    def _check_can_read(self):
+        if self._mode not in (_MODE_READ, _MODE_READ_EOF):
+            self._check_not_closed()
+            raise IOError("File not open for reading")
+
+    def _check_can_write(self):
+        if self._mode != _MODE_WRITE:
+            self._check_not_closed()
+            raise IOError("File not open for writing")
+
+    def _check_can_seek(self):
+        if self._mode not in (_MODE_READ, _MODE_READ_EOF):
+            self._check_not_closed()
+            raise IOError("Seeking only works on files opening for reading")
+
+    # Fill the readahead buffer if it is empty. Returns False on EOF.
+    def _fill_buffer(self):
+        if self._buffer:
+            return True
+        if self._decompressor.eof:
+            self._mode = _MODE_READ_EOF
+            self._size = self._pos
+            return False
+        rawblock = self._fp.read(_BUFFER_SIZE)
+        if not rawblock:
+            raise EOFError("Compressed file ended before the "
+                           "end-of-stream marker was reached")
+        self._buffer = self._decompressor.decompress(rawblock)
+        return True
+
+    # Read data until EOF.
+    # If return_data is false, consume the data without returning it.
+    def _read_all(self, return_data=True):
+        blocks = []
+        while self._fill_buffer():
+            if return_data:
+                blocks.append(self._buffer)
+            self._pos += len(self._buffer)
+            self._buffer = None
+        if return_data:
+            return b"".join(blocks)
+
+    # Read a block of up to size bytes.
+    # If return_data is false, consume the data without returning it.
+    def _read_block(self, size, return_data=True):
+        blocks = []
+        nread = 0
+        while nread < size and self._fill_buffer():
+            nbuf = len(self._buffer)
+            ct = min(nbuf, size - nread)
+            if return_data:
+                blocks.append(self._buffer[:ct] if ct < nbuf else self._buffer)
+            nread += ct
+            self._pos += ct
+            self._buffer = self._buffer[ct:] if ct < nbuf else None
+        if return_data:
+            return b"".join(blocks)
+
+    # Read a single line of up to size bytes. Negative size -> no size limit.
+    def _read_line(self, size):
+        blocks = []
+        nread = 0
+        while (size < 0 or nread < size) and self._fill_buffer():
+            nbuf = len(self._buffer)
+            if size < 0 or nbuf <= size - nread:
+                limit = nbuf
+            else:
+                limit = size - nread
+            try:
+                ct = self._buffer.index(b'\n', 0, limit) + 1
+                found_eol = True
+            except ValueError:
+                ct = limit
+                found_eol = False
+            blocks.append(self._buffer[:ct] if ct < nbuf else self._buffer)
+            self._pos += ct
+            self._buffer = self._buffer[ct:] if ct < nbuf else None
+            if found_eol:
+                break
+        return b"".join(blocks)
+
+    # Compress and write a block of data.
+    def _write(self, data):
+        compressed = self._compressor.compress(data)
+        self._fp.write(compressed)
+        self._pos += len(data)
+        return len(data)
+
+    # Rewind the file to the beginning of the data stream.
+    def _rewind(self):
+        self._fp.seek(0, 0)
+        self._mode = _MODE_READ
+        self._pos = 0
+        self._decompressor = _bz2.Decompressor()
+        self._buffer = None
+
+
+class BZ2Compressor:

[Antoine Pitrou, 2011/03/13 17:31:28]

I don't think there's a point in a Python wrapper, since the wrapper is so
trivial. Just do the lock operations in C.
Same for BZ2Decompressor.

+
+    """A compressor object for compressing data incrementally.
+
+    For one-shot compression, use the compress() function instead.
+    """
+
+    def __init__(self, compresslevel=9):
+        """Create a new compressor object.
+
+        compresslevel, if given, must be a number between 1 and 9.
+        """
+        if not (1 <= compresslevel <= 9):
+            raise ValueError("compresslevel must be between 1 and 9")
+        self._lock = threading.Lock()
+        self._compressor = _bz2.Compressor(compresslevel)
+
+    def compress(self, data):
+        """Provide data to the compressor object.
+
+        Returns a chunk of decompressed data if possible, or b''
+        otherwise.
+
+        When you have finished providing data to the compressor, call
+        the flush() method to finish the compression process.
+        """
+        with self._lock:
+            return self._compressor.compress(data)
+
+    def flush(self):
+        """Flush the internal buffers of the compressor object.
+
+        Returns any compressed data that was not returned by earlier
+        calls to decompress().
+
+        Once this method has been called, the compressor object may no
+        longer be used.
+        """
+        with self._lock:
+            return self._compressor.flush()
+
+
+class BZ2Decompressor:
+
+    """A decompressor object for decompressing data incrementally.
+
+    For one-shot decompression, use the decompress() function instead.
+    """
+
+    def __init__(self):
+        """Create a new decompressor object."""
+        self._lock = threading.Lock()
+        self._decompressor = _bz2.Decompressor()
+
+    def decompress(self, data):
+        """Provide data to the decompressor object.
+
+        Returns a chunk of decompressed data if possible, or b''
+        otherwise.
+
+        Attempting to decompress data after the end of stream is reached
+        raises an EOFError. Any data found after the end of the stream
+        is ignored and saved in the unused_data attribute.
+        """
+        with self._lock:
+            return self._decompressor.decompress(data)
+
+    @property
+    def eof(self):
+        """True if the end-of-stream marker has been reached."""
+        return self._decompressor.eof
+
+    @property
+    def unused_data(self):
+        """Data found after the end of the compressed stream."""
+        return self._decompressor.unused_data
+
+
+def compress(data, compresslevel=9):
+    """Compress a block of data.
+
+    compresslevel, if given, must be a number between 1 and 9.
+
+    For incremental compression, use a BZ2Compressor object instead.
+    """
+    comp = _bz2.Compressor(compresslevel)
+    return comp.compress(data) + comp.flush()
+
+
+def decompress(data):
+    """Decompress a block of data.
+
+    For incremental decompression, use a BZ2Decompressor object instead.
+    """
+    if len(data) == 0:
+        return b""
+    decomp = _bz2.Decompressor()
+    result = decomp.decompress(data)
+    if not decomp.eof:
+        raise ValueError("Compressed data ended before the "
+                         "end-of-stream marker was reached")
+    return result