Issue #19223: Add support for the 'x' mode to the bz2 module.

Patch by Tim Heaney and Vajrasky Kok.

2 files changed, 67 insertions, 42 deletions

diff --git a/Lib/bz2.py b/Lib/bz2.py
index 6e6a2b9948..6bc611e71f 100644
--- a/Lib/bz2.py
+++ b/Lib/bz2.py
@@ -49,12 +49,12 @@ class BZ2File(io.BufferedIOBase):
         which will be used to read or write the compressed data.
 
         mode can be 'r' for reading (default), 'w' for (over)writing,
-        or 'a' for appending. These can equivalently be given as 'rb',
-        'wb', and 'ab'.
+        'x' for creating exclusively, or 'a' for appending. These can
+        equivalently be given as 'rb', 'wb', 'xb', and 'ab'.
 
         buffering is ignored. Its use is deprecated.
 
-        If mode is 'w' or 'a', compresslevel can be a number between 1
+        If mode is 'w', 'x' or 'a', 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.
 
@@ -87,6 +87,10 @@ class BZ2File(io.BufferedIOBase):
             mode = "wb"
             mode_code = _MODE_WRITE
             self._compressor = BZ2Compressor(compresslevel)
+        elif mode in ("x", "xb"):
+            mode = "xb"
+            mode_code = _MODE_WRITE
+            self._compressor = BZ2Compressor(compresslevel)
         elif mode in ("a", "ab"):
             mode = "ab"
             mode_code = _MODE_WRITE
@@ -443,9 +447,9 @@ def open(filename, mode="rb", compresslevel=9,
     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.
+    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.
 
     For binary mode, this function is equivalent to the BZ2File
     constructor: BZ2File(filename, mode, compresslevel). In this case,
diff --git a/Lib/test/test_bz2.py b/Lib/test/test_bz2.py
index 7090cd6935..d087cc3ee9 100644
--- a/Lib/test/test_bz2.py
+++ b/Lib/test/test_bz2.py
@@ -8,6 +8,7 @@ import os
 import random
 import subprocess
 import sys
+from test.support import unlink
 
 try:
     import threading
@@ -715,49 +716,67 @@ class OpenTest(BaseTest):
         return bz2.open(*args, **kwargs)
 
     def test_binary_modes(self):
-        with self.open(self.filename, "wb") as f:
-            f.write(self.TEXT)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read())
-            self.assertEqual(file_data, self.TEXT)
-        with self.open(self.filename, "rb") as f:
-            self.assertEqual(f.read(), self.TEXT)
-        with self.open(self.filename, "ab") as f:
-            f.write(self.TEXT)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read())
-            self.assertEqual(file_data, self.TEXT * 2)
+        for mode in ("wb", "xb"):
+            if mode == "xb":
+                unlink(self.filename)
+            with self.open(self.filename, mode) as f:
+                f.write(self.TEXT)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read())
+                self.assertEqual(file_data, self.TEXT)
+            with self.open(self.filename, "rb") as f:
+                self.assertEqual(f.read(), self.TEXT)
+            with self.open(self.filename, "ab") as f:
+                f.write(self.TEXT)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read())
+                self.assertEqual(file_data, self.TEXT * 2)
 
     def test_implicit_binary_modes(self):
         # Test implicit binary modes (no "b" or "t" in mode string).
-        with self.open(self.filename, "w") as f:
-            f.write(self.TEXT)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read())
-            self.assertEqual(file_data, self.TEXT)
-        with self.open(self.filename, "r") as f:
-            self.assertEqual(f.read(), self.TEXT)
-        with self.open(self.filename, "a") as f:
-            f.write(self.TEXT)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read())
-            self.assertEqual(file_data, self.TEXT * 2)
+        for mode in ("w", "x"):
+            if mode == "x":
+                unlink(self.filename)
+            with self.open(self.filename, mode) as f:
+                f.write(self.TEXT)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read())
+                self.assertEqual(file_data, self.TEXT)
+            with self.open(self.filename, "r") as f:
+                self.assertEqual(f.read(), self.TEXT)
+            with self.open(self.filename, "a") as f:
+                f.write(self.TEXT)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read())
+                self.assertEqual(file_data, self.TEXT * 2)
 
     def test_text_modes(self):
         text = self.TEXT.decode("ascii")
         text_native_eol = text.replace("\n", os.linesep)
-        with self.open(self.filename, "wt") as f:
-            f.write(text)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read()).decode("ascii")
-            self.assertEqual(file_data, text_native_eol)
-        with self.open(self.filename, "rt") as f:
-            self.assertEqual(f.read(), text)
-        with self.open(self.filename, "at") as f:
-            f.write(text)
-        with open(self.filename, "rb") as f:
-            file_data = self.decompress(f.read()).decode("ascii")
-            self.assertEqual(file_data, text_native_eol * 2)
+        for mode in ("wt", "xt"):
+            if mode == "xt":
+                unlink(self.filename)
+            with self.open(self.filename, mode) as f:
+                f.write(text)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read()).decode("ascii")
+                self.assertEqual(file_data, text_native_eol)
+            with self.open(self.filename, "rt") as f:
+                self.assertEqual(f.read(), text)
+            with self.open(self.filename, "at") as f:
+                f.write(text)
+            with open(self.filename, "rb") as f:
+                file_data = self.decompress(f.read()).decode("ascii")
+                self.assertEqual(file_data, text_native_eol * 2)
+
+    def test_x_mode(self):
+        for mode in ("x", "xb", "xt"):
+            unlink(self.filename)
+            with self.open(self.filename, mode) as f:
+                pass
+            with self.assertRaises(FileExistsError):
+                with self.open(self.filename, mode) as f:
+                    pass
 
     def test_fileobj(self):
         with self.open(BytesIO(self.DATA), "r") as f:
@@ -773,6 +792,8 @@ class OpenTest(BaseTest):
         self.assertRaises(ValueError,
                           self.open, self.filename, "wbt")
         self.assertRaises(ValueError,
+                          self.open, self.filename, "xbt")
+        self.assertRaises(ValueError,
                           self.open, self.filename, "rb", encoding="utf-8")
         self.assertRaises(ValueError,
                           self.open, self.filename, "rb", errors="ignore")