Add a function bz2.open(), to match gzip.open().

1 files changed, 44 insertions, 3 deletions

diff --git a/Lib/bz2.py b/Lib/bz2.py
index d04731b..2e808cd 100644
--- a/Lib/bz2.py
+++ b/Lib/bz2.py
@@ -4,11 +4,12 @@ This module provides a file interface, classes for incremental
 (de)compression, and functions for one-shot (de)compression.
 """
 
-__all__ = ["BZ2File", "BZ2Compressor", "BZ2Decompressor", "compress",
-           "decompress"]
+__all__ = ["BZ2File", "BZ2Compressor", "BZ2Decompressor",
+           "open", "compress", "decompress"]
 
 __author__ = "Nadeem Vawda <nadeem.vawda@gmail.com>"
 
+import builtins
 import io
 import warnings
 
@@ -91,7 +92,7 @@ class BZ2File(io.BufferedIOBase):
             raise ValueError("Invalid mode: {!r}".format(mode))
 
         if isinstance(filename, (str, bytes)):
-            self._fp = open(filename, mode)
+            self._fp = builtins.open(filename, mode)
             self._closefp = True
             self._mode = mode_code
         elif hasattr(filename, "read") or hasattr(filename, "write"):
@@ -391,6 +392,46 @@ class BZ2File(io.BufferedIOBase):
             return self._pos
 
 
+def open(filename, mode="rb", compresslevel=9,
+         encoding=None, errors=None, newline=None):
+    """Open a bzip2-compressed file in binary or text mode.
+
+    The filename argument can be an actual filename (a str or bytes object), or
+    an existing file object to read from or write to.
+
+    The mode argument can be "r", "rb", "w", "wb", "a" or "ab" for binary mode,
+    or "rt", "wt" or "at" for text mode. The default mode is "rb", and the
+    default compresslevel is 9.
+
+    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.
+
+    For text mode, a BZ2File object is created, and wrapped in an
+    io.TextIOWrapper instance with the specified encoding, error handling
+    behavior, and line ending(s).
+
+    """
+    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
+
+
 def compress(data, compresslevel=9):
     """Compress a block of data.