Allow LZMAFile to accept modes with a "b" suffix.

4 files changed, 30 insertions, 9 deletions

diff --git a/Doc/library/lzma.rst b/Doc/library/lzma.rst
index 156d77a..67e425d 100644
--- a/Doc/library/lzma.rst
+++ b/Doc/library/lzma.rst
@@ -40,9 +40,11 @@ Reading and writing compressed files
    file will not be closed when the :class:`LZMAFile` is closed.
 
    The *mode* argument can be either ``"r"`` for reading (default), ``"w"`` for
-   overwriting, or ``"a"`` for appending. If *filename* is an existing file
-   object, a mode of ``"w"`` does not truncate the file, and is instead
-   equivalent to ``"a"``.
+   overwriting, or ``"a"`` for appending. These can equivalently be given as
+   ``"rb"``, ``"wb"``, and ``"ab"`` respectively.
+
+   If *filename* is a file object (rather than an actual file name), a mode of
+   ``"w"`` does not truncate the file, and is instead equivalent to ``"a"``.
 
    When opening a file for reading, the input file may be the concatenation of
    multiple separate compressed streams. These are transparently decoded as a
diff --git a/Lib/lzma.py b/Lib/lzma.py
index 8fb3f03..0790691 100644
--- a/Lib/lzma.py
+++ b/Lib/lzma.py
@@ -54,7 +54,8 @@ class LZMAFile(io.BufferedIOBase):
         be an existing file object to read from or write to.
 
         mode can be "r" for reading (default), "w" for (over)writing, or
-        "a" for appending.
+        "a" for appending. These can equivalently be given as "rb", "wb",
+        and "ab" respectively.
 
         format specifies the container format to use for the file.
         If mode is "r", this defaults to FORMAT_AUTO. Otherwise, the
@@ -93,7 +94,7 @@ class LZMAFile(io.BufferedIOBase):
         self._pos = 0
         self._size = -1
 
-        if mode == "r":
+        if mode in ("r", "rb"):
             if check != -1:
                 raise ValueError("Cannot specify an integrity check "
                                  "when opening a file for reading")
@@ -109,7 +110,7 @@ class LZMAFile(io.BufferedIOBase):
             self._init_args = {"format":format, "filters":filters}
             self._decompressor = LZMADecompressor(**self._init_args)
             self._buffer = None
-        elif mode in ("w", "a"):
+        elif mode in ("w", "wb", "a", "ab"):
             if format is None:
                 format = FORMAT_XZ
             mode_code = _MODE_WRITE
@@ -119,7 +120,8 @@ class LZMAFile(io.BufferedIOBase):
             raise ValueError("Invalid mode: {!r}".format(mode))
 
         if isinstance(filename, (str, bytes)):
-            mode += "b"
+            if "b" not in mode:
+                mode += "b"
             self._fp = open(filename, mode)
             self._closefp = True
             self._mode = mode_code
diff --git a/Lib/test/test_lzma.py b/Lib/test/test_lzma.py
index aee7921..e4d2cb1 100644
--- a/Lib/test/test_lzma.py
+++ b/Lib/test/test_lzma.py
@@ -374,6 +374,21 @@ class FileTestCase(unittest.TestCase):
             with LZMAFile(TESTFN, "a") as f:
                 pass
 
+    def test_init_mode(self):
+        with TempFile(TESTFN):
+            with LZMAFile(TESTFN, "r"):
+                pass
+            with LZMAFile(TESTFN, "rb"):
+                pass
+            with LZMAFile(TESTFN, "w"):
+                pass
+            with LZMAFile(TESTFN, "wb"):
+                pass
+            with LZMAFile(TESTFN, "a"):
+                pass
+            with LZMAFile(TESTFN, "ab"):
+                pass
+
     def test_init_bad_mode(self):
         with self.assertRaises(ValueError):
             LZMAFile(BytesIO(COMPRESSED_XZ), (3, "x"))
@@ -382,11 +397,11 @@ class FileTestCase(unittest.TestCase):
         with self.assertRaises(ValueError):
             LZMAFile(BytesIO(COMPRESSED_XZ), "x")
         with self.assertRaises(ValueError):
-            LZMAFile(BytesIO(COMPRESSED_XZ), "rb")
+            LZMAFile(BytesIO(COMPRESSED_XZ), "rt")
         with self.assertRaises(ValueError):
             LZMAFile(BytesIO(COMPRESSED_XZ), "r+")
         with self.assertRaises(ValueError):
-            LZMAFile(BytesIO(COMPRESSED_XZ), "wb")
+            LZMAFile(BytesIO(COMPRESSED_XZ), "wt")
         with self.assertRaises(ValueError):
             LZMAFile(BytesIO(COMPRESSED_XZ), "w+")
         with self.assertRaises(ValueError):
diff --git a/Misc/NEWS b/Misc/NEWS
index 24eeeaa..333999f 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -15,6 +15,8 @@ Core and Builtins
 Library
 -------
 
+- LZMAFile now accepts the modes "rb"/"wb"/"ab" as synonyms of "r"/"w"/"a".
+
 - The bz2 module now contains an open() function, allowing compressed files to
   conveniently be opened in text mode as well as binary mode.