bz2.py 12.3 KB
Newer Older
1 2 3 4 5 6
"""Interface to the libbzip2 compression library.

This module provides a file interface, classes for incremental
(de)compression, and functions for one-shot (de)compression.
"""

7 8
__all__ = ["BZ2File", "BZ2Compressor", "BZ2Decompressor",
           "open", "compress", "decompress"]
9 10 11

__author__ = "Nadeem Vawda <nadeem.vawda@gmail.com>"

12
from builtins import open as _builtin_open
13
import io
14
import os
15
import warnings
16
import _compression
17
from threading import RLock
18

19 20 21 22 23
from _bz2 import BZ2Compressor, BZ2Decompressor


_MODE_CLOSED   = 0
_MODE_READ     = 1
24
# Value 2 no longer used
25 26
_MODE_WRITE    = 3

27 28
_sentinel = object()

29

30
class BZ2File(_compression.BaseStream):
31 32 33 34 35 36 37 38 39 40

    """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.
    """

41
    def __init__(self, filename, mode="r", buffering=_sentinel, compresslevel=9):
42 43
        """Open a bzip2-compressed file.

44 45 46
        If filename is a str, bytes, or PathLike object, it gives the
        name of the file to be opened. Otherwise, it should be a file
        object, which will be used to read or write the compressed data.
47

48
        mode can be 'r' for reading (default), 'w' for (over)writing,
49 50
        'x' for creating exclusively, or 'a' for appending. These can
        equivalently be given as 'rb', 'wb', 'xb', and 'ab'.
51

52
        buffering is ignored since Python 3.0. Its use is deprecated.
53

54
        If mode is 'w', 'x' or 'a', compresslevel can be a number between 1
55
        and 9 specifying the level of compression: 1 produces the least
56
        compression, and 9 (default) produces the most compression.
57 58 59

        If mode is 'r', the input file may be the concatenation of
        multiple compressed streams.
60 61
        """
        # This lock must be recursive, so that BufferedIOBase's
62
        # writelines() does not deadlock.
63
        self._lock = RLock()
64 65 66 67
        self._fp = None
        self._closefp = False
        self._mode = _MODE_CLOSED

68
        if buffering is not _sentinel:
69
            warnings.warn("Use of 'buffering' argument is deprecated and ignored "
70 71 72
                          "since Python 3.0.",
                          DeprecationWarning,
                          stacklevel=2)
73 74 75 76 77 78 79 80 81 82

        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
        elif mode in ("w", "wb"):
            mode = "wb"
            mode_code = _MODE_WRITE
83
            self._compressor = BZ2Compressor(compresslevel)
84 85 86 87
        elif mode in ("x", "xb"):
            mode = "xb"
            mode_code = _MODE_WRITE
            self._compressor = BZ2Compressor(compresslevel)
88 89 90
        elif mode in ("a", "ab"):
            mode = "ab"
            mode_code = _MODE_WRITE
91
            self._compressor = BZ2Compressor(compresslevel)
92
        else:
93
            raise ValueError("Invalid mode: %r" % (mode,))
94

95
        if isinstance(filename, (str, bytes, os.PathLike)):
96
            self._fp = _builtin_open(filename, mode)
97 98
            self._closefp = True
            self._mode = mode_code
99 100
        elif hasattr(filename, "read") or hasattr(filename, "write"):
            self._fp = filename
101 102
            self._mode = mode_code
        else:
103
            raise TypeError("filename must be a str, bytes, file or PathLike object")
104

105 106 107 108 109 110 111
        if self._mode == _MODE_READ:
            raw = _compression.DecompressReader(self._fp,
                BZ2Decompressor, trailing_error=OSError)
            self._buffer = io.BufferedReader(raw)
        else:
            self._pos = 0

112 113 114 115 116 117 118 119 120 121
    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:
            if self._mode == _MODE_CLOSED:
                return
            try:
122 123
                if self._mode == _MODE_READ:
                    self._buffer.close()
124 125 126 127
                elif self._mode == _MODE_WRITE:
                    self._fp.write(self._compressor.flush())
                    self._compressor = None
            finally:
Antoine Pitrou's avatar
Antoine Pitrou committed
128
                try:
129 130 131 132 133 134
                    if self._closefp:
                        self._fp.close()
                finally:
                    self._fp = None
                    self._closefp = False
                    self._mode = _MODE_CLOSED
135
                    self._buffer = None
136 137 138 139 140 141 142 143

    @property
    def closed(self):
        """True if this file is closed."""
        return self._mode == _MODE_CLOSED

    def fileno(self):
        """Return the file descriptor for the underlying file."""
144
        self._check_not_closed()
145 146 147 148
        return self._fp.fileno()

    def seekable(self):
        """Return whether the file supports seeking."""
149
        return self.readable() and self._buffer.seekable()
150 151 152

    def readable(self):
        """Return whether the file was opened for reading."""
153
        self._check_not_closed()
154
        return self._mode == _MODE_READ
155 156 157

    def writable(self):
        """Return whether the file was opened for writing."""
158
        self._check_not_closed()
159 160 161 162 163 164 165 166 167 168
        return self._mode == _MODE_WRITE

    def peek(self, n=0):
        """Return buffered data without advancing the file position.

        Always returns at least one byte of data, unless at EOF.
        The exact number of bytes returned is unspecified.
        """
        with self._lock:
            self._check_can_read()
169 170 171 172
            # Relies on the undocumented fact that BufferedReader.peek()
            # always returns at least one byte (except at EOF), independent
            # of the value of n
            return self._buffer.peek(n)
173 174 175 176 177 178 179 180 181

    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()
182
            return self._buffer.read(size)
183 184

    def read1(self, size=-1):
185
        """Read up to size uncompressed bytes, while trying to avoid
186 187
        making multiple reads from the underlying stream. Reads up to a
        buffer's worth of data if size is negative.
188 189 190 191 192

        Returns b'' if the file is at EOF.
        """
        with self._lock:
            self._check_can_read()
193 194 195
            if size < 0:
                size = io.DEFAULT_BUFFER_SIZE
            return self._buffer.read1(size)
196 197

    def readinto(self, b):
198
        """Read bytes into b.
Antoine Pitrou's avatar
Antoine Pitrou committed
199

200 201 202
        Returns the number of bytes read (0 for EOF).
        """
        with self._lock:
203 204
            self._check_can_read()
            return self._buffer.readinto(b)
205 206 207 208 209 210 211 212

    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.
        """
213 214 215 216
        if not isinstance(size, int):
            if not hasattr(size, "__index__"):
                raise TypeError("Integer argument expected")
            size = size.__index__()
217
        with self._lock:
218
            self._check_can_read()
219
            return self._buffer.readline(size)
220 221 222 223 224 225 226 227

    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.
        """
228 229 230 231
        if not isinstance(size, int):
            if not hasattr(size, "__index__"):
                raise TypeError("Integer argument expected")
            size = size.__index__()
232
        with self._lock:
233 234
            self._check_can_read()
            return self._buffer.readlines(size)
235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258

    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()
            compressed = self._compressor.compress(data)
            self._fp.write(compressed)
            self._pos += len(data)
            return len(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:
259 260 261
            return _compression.BaseStream.writelines(self, seq)

    def seek(self, offset, whence=io.SEEK_SET):
262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277
        """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()
278
            return self._buffer.seek(offset, whence)
279 280 281 282 283

    def tell(self):
        """Return the current file position."""
        with self._lock:
            self._check_not_closed()
284 285
            if self._mode == _MODE_READ:
                return self._buffer.tell()
286 287 288
            return self._pos


289 290 291 292
def open(filename, mode="rb", compresslevel=9,
         encoding=None, errors=None, newline=None):
    """Open a bzip2-compressed file in binary or text mode.

293 294 295
    The filename argument can be an actual filename (a str, bytes, or
    PathLike object), or an existing file object to read from or write
    to.
296

297 298 299
    The mode argument can be "r", "rb", "w", "wb", "x", "xb", "a" or
    "ab" for binary mode, or "rt", "wt", "xt" or "at" for text mode.
    The default mode is "rb", and the default compresslevel is 9.
300

301 302 303
    For binary mode, this function is equivalent to the BZ2File
    constructor: BZ2File(filename, mode, compresslevel). In this case,
    the encoding, errors and newline arguments must not be provided.
304 305

    For text mode, a BZ2File object is created, and wrapped in an
306 307
    io.TextIOWrapper instance with the specified encoding, error
    handling behavior, and line ending(s).
308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329

    """
    if "t" in mode:
        if "b" in mode:
            raise ValueError("Invalid mode: %r" % (mode,))
    else:
        if encoding is not None:
            raise ValueError("Argument 'encoding' not supported in binary mode")
        if errors is not None:
            raise ValueError("Argument 'errors' not supported in binary mode")
        if newline is not None:
            raise ValueError("Argument 'newline' not supported in binary mode")

    bz_mode = mode.replace("t", "")
    binary_file = BZ2File(filename, bz_mode, compresslevel=compresslevel)

    if "t" in mode:
        return io.TextIOWrapper(binary_file, encoding, errors, newline)
    else:
        return binary_file


330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345
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 = BZ2Compressor(compresslevel)
    return comp.compress(data) + comp.flush()


def decompress(data):
    """Decompress a block of data.

    For incremental decompression, use a BZ2Decompressor object instead.
    """
346
    results = []
347
    while data:
348
        decomp = BZ2Decompressor()
349 350 351 352 353 354 355 356
        try:
            res = decomp.decompress(data)
        except OSError:
            if results:
                break  # Leftover data is not a valid bzip2 stream; ignore it.
            else:
                raise  # Error on the first iteration; bail out.
        results.append(res)
357 358 359 360
        if not decomp.eof:
            raise ValueError("Compressed data ended before the "
                             "end-of-stream marker was reached")
        data = decomp.unused_data
361
    return b"".join(results)