bpo-14156: Make argparse.FileType work correctly for binary file modes when argument is '-' (GH-13165)

miss-islington · MojoVampire · web-flow · commit ee18df425209 · 2022-03-06T04:12:06.000-08:00
Also made modes containing 'a' or 'x' act the same as a mode containing 'w' when argument is '-' (so 'a'/'x' return sys.stdout like 'w', and 'ab'/'xb' return sys.stdout.buffer like 'wb'). (cherry picked from commit eafec26) Co-authored-by: MojoVampire <shadowranger+github@gmail.com>
diff --git a/Lib/argparse.py b/Lib/argparse.py
@@ -726,7 +726,7 @@ def _get_action_name(argument):
     if argument is None:
         return None
     elif argument.option_strings:
-        return  '/'.join(argument.option_strings)
+        return '/'.join(argument.option_strings)
     elif argument.metavar not in (None, SUPPRESS):
         return argument.metavar
     elif argument.dest not in (None, SUPPRESS):
@@ -1256,9 +1256,9 @@ def __call__(self, string):
         # the special argument "-" means sys.std{in,out}
         if string == '-':
             if 'r' in self._mode:
-                return _sys.stdin
-            elif 'w' in self._mode:
-                return _sys.stdout
+                return _sys.stdin.buffer if 'b' in self._mode else _sys.stdin
+            elif any(c in self._mode for c in 'wax'):
+                return _sys.stdout.buffer if 'b' in self._mode else _sys.stdout
             else:
                 msg = _('argument "-" with mode %r') % self._mode
                 raise ValueError(msg)
diff --git a/Lib/test/test_argparse.py b/Lib/test/test_argparse.py
@@ -1,6 +1,8 @@
 # Author: Steven J. Bethard <steven.bethard@gmail.com>.
 
 import inspect
+import io
+import operator
 import os
 import shutil
 import stat
@@ -10,12 +12,27 @@
 import unittest
 import argparse
 
-from io import StringIO
-
 from test.support import os_helper
 from unittest import mock
-class StdIOBuffer(StringIO):
-    pass
+
+
+class StdIOBuffer(io.TextIOWrapper):
+    '''Replacement for writable io.StringIO that behaves more like real file
+
+    Unlike StringIO, provides a buffer attribute that holds the underlying
+    binary data, allowing it to replace sys.stdout/sys.stderr in more
+    contexts.
+    '''
+
+    def __init__(self, initial_value='', newline='\n'):
+        initial_value = initial_value.encode('utf-8')
+        super().__init__(io.BufferedWriter(io.BytesIO(initial_value)),
+                         'utf-8', newline=newline)
+
+    def getvalue(self):
+        self.flush()
+        return self.buffer.raw.getvalue().decode('utf-8')
+
 
 class TestCase(unittest.TestCase):
 
@@ -42,11 +59,14 @@ def tearDown(self):
                 os.chmod(os.path.join(self.temp_dir, name), stat.S_IWRITE)
         shutil.rmtree(self.temp_dir, True)
 
-    def create_readonly_file(self, filename):
+    def create_writable_file(self, filename):
         file_path = os.path.join(self.temp_dir, filename)
         with open(file_path, 'w', encoding="utf-8") as file:
             file.write(filename)
-        os.chmod(file_path, stat.S_IREAD)
+        return file_path
+
+    def create_readonly_file(self, filename):
+        os.chmod(self.create_writable_file(filename), stat.S_IREAD)
 
 class Sig(object):
 
@@ -96,10 +116,15 @@ def stderr_to_parser_error(parse_args, *args, **kwargs):
         try:
             result = parse_args(*args, **kwargs)
             for key in list(vars(result)):
-                if getattr(result, key) is sys.stdout:
+                attr = getattr(result, key)
+                if attr is sys.stdout:
                     setattr(result, key, old_stdout)
-                if getattr(result, key) is sys.stderr:
+                elif attr is sys.stdout.buffer:
+                    setattr(result, key, getattr(old_stdout, 'buffer', BIN_STDOUT_SENTINEL))
+                elif attr is sys.stderr:
                     setattr(result, key, old_stderr)
+                elif attr is sys.stderr.buffer:
+                    setattr(result, key, getattr(old_stderr, 'buffer', BIN_STDERR_SENTINEL))
             return result
         except SystemExit as e:
             code = e.code
@@ -1545,16 +1570,40 @@ def test_r_1_replace(self):
         type = argparse.FileType('r', 1, errors='replace')
         self.assertEqual("FileType('r', 1, errors='replace')", repr(type))
 
+
+BIN_STDOUT_SENTINEL = object()
+BIN_STDERR_SENTINEL = object()
+
+
 class StdStreamComparer:
     def __init__(self, attr):
-        self.attr = attr
+        # We try to use the actual stdXXX.buffer attribute as our
+        # marker, but but under some test environments,
+        # sys.stdout/err are replaced by io.StringIO which won't have .buffer,
+        # so we use a sentinel simply to show that the tests do the right thing
+        # for any buffer supporting object
+        self.getattr = operator.attrgetter(attr)
+        if attr == 'stdout.buffer':
+            self.backupattr = BIN_STDOUT_SENTINEL
+        elif attr == 'stderr.buffer':
+            self.backupattr = BIN_STDERR_SENTINEL
+        else:
+            self.backupattr = object() # Not equal to anything
 
     def __eq__(self, other):
-        return other == getattr(sys, self.attr)
+        try:
+            return other == self.getattr(sys)
+        except AttributeError:
+            return other == self.backupattr
+
 
 eq_stdin = StdStreamComparer('stdin')
 eq_stdout = StdStreamComparer('stdout')
 eq_stderr = StdStreamComparer('stderr')
+eq_bstdin = StdStreamComparer('stdin.buffer')
+eq_bstdout = StdStreamComparer('stdout.buffer')
+eq_bstderr = StdStreamComparer('stderr.buffer')
+
 
 class RFile(object):
     seen = {}
@@ -1633,7 +1682,7 @@ def setUp(self):
         ('foo', NS(x=None, spam=RFile('foo'))),
         ('-x foo bar', NS(x=RFile('foo'), spam=RFile('bar'))),
         ('bar -x foo', NS(x=RFile('foo'), spam=RFile('bar'))),
-        ('-x - -', NS(x=eq_stdin, spam=eq_stdin)),
+        ('-x - -', NS(x=eq_bstdin, spam=eq_bstdin)),
     ]
 
 
@@ -1660,8 +1709,9 @@ class TestFileTypeW(TempDirMixin, ParserTestCase):
     """Test the FileType option/argument type for writing files"""
 
     def setUp(self):
-        super(TestFileTypeW, self).setUp()
+        super().setUp()
         self.create_readonly_file('readonly')
+        self.create_writable_file('writable')
 
     argument_signatures = [
         Sig('-x', type=argparse.FileType('w')),
@@ -1670,13 +1720,37 @@ def setUp(self):
     failures = ['-x', '', 'readonly']
     successes = [
         ('foo', NS(x=None, spam=WFile('foo'))),
+        ('writable', NS(x=None, spam=WFile('writable'))),
         ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),
         ('bar -x foo', NS(x=WFile('foo'), spam=WFile('bar'))),
         ('-x - -', NS(x=eq_stdout, spam=eq_stdout)),
     ]
 
+@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,
+                 "non-root user required")
+class TestFileTypeX(TempDirMixin, ParserTestCase):
+    """Test the FileType option/argument type for writing new files only"""
+
+    def setUp(self):
+        super().setUp()
+        self.create_readonly_file('readonly')
+        self.create_writable_file('writable')
+
+    argument_signatures = [
+        Sig('-x', type=argparse.FileType('x')),
+        Sig('spam', type=argparse.FileType('x')),
+    ]
+    failures = ['-x', '', 'readonly', 'writable']
+    successes = [
+        ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),
+        ('-x - -', NS(x=eq_stdout, spam=eq_stdout)),
+    ]
+
 
+@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,
+                 "non-root user required")
 class TestFileTypeWB(TempDirMixin, ParserTestCase):
+    """Test the FileType option/argument type for writing binary files"""
 
     argument_signatures = [
         Sig('-x', type=argparse.FileType('wb')),
@@ -1687,7 +1761,22 @@ class TestFileTypeWB(TempDirMixin, ParserTestCase):
         ('foo', NS(x=None, spam=WFile('foo'))),
         ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),
         ('bar -x foo', NS(x=WFile('foo'), spam=WFile('bar'))),
-        ('-x - -', NS(x=eq_stdout, spam=eq_stdout)),
+        ('-x - -', NS(x=eq_bstdout, spam=eq_bstdout)),
+    ]
+
+
+@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,
+                 "non-root user required")
+class TestFileTypeXB(TestFileTypeX):
+    "Test the FileType option/argument type for writing new binary files only"
+
+    argument_signatures = [
+        Sig('-x', type=argparse.FileType('xb')),
+        Sig('spam', type=argparse.FileType('xb')),
+    ]
+    successes = [
+        ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),
+        ('-x - -', NS(x=eq_bstdout, spam=eq_bstdout)),
     ]
 
 
diff --git a/Misc/NEWS.d/next/Library/2019-05-07-14-25-45.bpo-14156.0FaHXE.rst b/Misc/NEWS.d/next/Library/2019-05-07-14-25-45.bpo-14156.0FaHXE.rst
@@ -0,0 +1,4 @@
+argparse.FileType now supports an argument of '-' in binary mode, returning
+the .buffer attribute of sys.stdin/sys.stdout as appropriate. Modes
+including 'x' and 'a' are treated equivalently to 'w' when argument is '-'.
+Patch contributed by Josh Rosenberg
