diff options
Diffstat (limited to 'Lib/test/test_argparse.py')
| -rw-r--r-- | Lib/test/test_argparse.py | 115 |
1 files changed, 102 insertions, 13 deletions
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 <steven.bethard@gmail.com>. 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)), ] |