From eafec26ae5327bb23b6dace2650b074c3327dfa0 Mon Sep 17 00:00:00 2001 From: MojoVampire Date: Sun, 6 Mar 2022 11:49:42 +0000 Subject: bpo-14156: Make argparse.FileType work correctly for binary file modes when argument is '-' (GH-13165) 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'). --- Lib/argparse.py | 8 +- Lib/test/test_argparse.py | 115 ++++++++++++++++++--- .../2019-05-07-14-25-45.bpo-14156.0FaHXE.rst | 4 + 3 files changed, 110 insertions(+), 17 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2019-05-07-14-25-45.bpo-14156.0FaHXE.rst diff --git a/Lib/argparse.py b/Lib/argparse.py index 3c6aa3c..429a72a 100644 --- a/Lib/argparse.py +++ b/Lib/argparse.py @@ -728,7 +728,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): @@ -1259,9 +1259,9 @@ class FileType(object): # 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 index df6da92..1f03b7f 100644 --- a/Lib/test/test_argparse.py +++ b/Lib/test/test_argparse.py @@ -1,6 +1,8 @@ # Author: Steven J. Bethard . import inspect +import io +import operator import os import shutil import stat @@ -11,12 +13,27 @@ import unittest import argparse import warnings -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): @@ -43,11 +60,14 @@ class TempDirMixin(object): 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): @@ -97,10 +117,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 @@ -1565,16 +1590,40 @@ class TestFileTypeRepr(TestCase): 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 = {} @@ -1653,7 +1702,7 @@ class TestFileTypeRB(TempDirMixin, ParserTestCase): ('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)), ] @@ -1680,8 +1729,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')), @@ -1690,13 +1740,37 @@ class TestFileTypeW(TempDirMixin, ParserTestCase): 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')), @@ -1707,7 +1781,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 new file mode 100644 index 0000000..7bfc917 --- /dev/null +++ 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 -- cgit v0.12