[3.9] bpo-14156: Make argparse.FileType work correctly for binary fil… · python/cpython@4d2099f (original) (raw)

1

`# Author: Steven J. Bethard steven.bethard@gmail.com.

2

3

`import inspect

4

import io

5

import operator

4

6

`import os

5

7

`import shutil

6

8

`import stat

`@@ -10,12 +12,27 @@

10

12

`import unittest

11

13

`import argparse

12

14

13

from io import StringIO

14

-

15

`from test import support

16

`from unittest import mock

17

class StdIOBuffer(StringIO):

18

pass

17

+

18

+

19

class StdIOBuffer(io.TextIOWrapper):

20

'''Replacement for writable io.StringIO that behaves more like real file

21

+

22

Unlike StringIO, provides a buffer attribute that holds the underlying

23

binary data, allowing it to replace sys.stdout/sys.stderr in more

24

contexts.

25

'''

26

+

27

def init(self, initial_value='', newline='\n'):

28

initial_value = initial_value.encode('utf-8')

29

super().init(io.BufferedWriter(io.BytesIO(initial_value)),

30

'utf-8', newline=newline)

31

+

32

def getvalue(self):

33

self.flush()

34

return self.buffer.raw.getvalue().decode('utf-8')

35

+

19

36

20

37

`class TestCase(unittest.TestCase):

21

38

`@@ -42,11 +59,14 @@ def tearDown(self):

42

59

`os.chmod(os.path.join(self.temp_dir, name), stat.S_IWRITE)

43

60

`shutil.rmtree(self.temp_dir, True)

44

61

45

def create_readonly_file(self, filename):

62

def create_writable_file(self, filename):

46

63

`file_path = os.path.join(self.temp_dir, filename)

47

64

`with open(file_path, 'w') as file:

48

65

`file.write(filename)

49

os.chmod(file_path, stat.S_IREAD)

66

return file_path

67

+

68

def create_readonly_file(self, filename):

69

os.chmod(self.create_writable_file(filename), stat.S_IREAD)

50

70

51

71

`class Sig(object):

52

72

`@@ -96,10 +116,15 @@ def stderr_to_parser_error(parse_args, *args, **kwargs):

96

116

`try:

97

117

`result = parse_args(*args, **kwargs)

98

118

`for key in list(vars(result)):

99

if getattr(result, key) is sys.stdout:

119

attr = getattr(result, key)

120

if attr is sys.stdout:

100

121

`setattr(result, key, old_stdout)

101

if getattr(result, key) is sys.stderr:

122

elif attr is sys.stdout.buffer:

123

setattr(result, key, getattr(old_stdout, 'buffer', BIN_STDOUT_SENTINEL))

124

elif attr is sys.stderr:

102

125

`setattr(result, key, old_stderr)

126

elif attr is sys.stderr.buffer:

127

setattr(result, key, getattr(old_stderr, 'buffer', BIN_STDERR_SENTINEL))

103

128

`return result

104

129

`except SystemExit as e:

105

130

`code = e.code

`@@ -1545,16 +1570,40 @@ def test_r_1_replace(self):

1545

1570

`type = argparse.FileType('r', 1, errors='replace')

1546

1571

`self.assertEqual("FileType('r', 1, errors='replace')", repr(type))

1547

1572

1573

+

1574

BIN_STDOUT_SENTINEL = object()

1575

BIN_STDERR_SENTINEL = object()

1576

+

1577

+

1548

1578

`class StdStreamComparer:

1549

1579

`def init(self, attr):

1550

self.attr = attr

1580

We try to use the actual stdXXX.buffer attribute as our

1581

marker, but but under some test environments,

1582

sys.stdout/err are replaced by io.StringIO which won't have .buffer,

1583

so we use a sentinel simply to show that the tests do the right thing

1584

for any buffer supporting object

1585

self.getattr = operator.attrgetter(attr)

1586

if attr == 'stdout.buffer':

1587

self.backupattr = BIN_STDOUT_SENTINEL

1588

elif attr == 'stderr.buffer':

1589

self.backupattr = BIN_STDERR_SENTINEL

1590

else:

1591

self.backupattr = object() # Not equal to anything

1551

1592

1552

1593

`def eq(self, other):

1553

return other == getattr(sys, self.attr)

1594

try:

1595

return other == self.getattr(sys)

1596

except AttributeError:

1597

return other == self.backupattr

1598

+

1554

1599

1555

1600

`eq_stdin = StdStreamComparer('stdin')

1556

1601

`eq_stdout = StdStreamComparer('stdout')

1557

1602

`eq_stderr = StdStreamComparer('stderr')

1603

eq_bstdin = StdStreamComparer('stdin.buffer')

1604

eq_bstdout = StdStreamComparer('stdout.buffer')

1605

eq_bstderr = StdStreamComparer('stderr.buffer')

1606

+

1558

1607

1559

1608

`class RFile(object):

1560

1609

`seen = {}

`@@ -1631,7 +1680,7 @@ def setUp(self):

1631

1680

` ('foo', NS(x=None, spam=RFile('foo'))),

1632

1681

` ('-x foo bar', NS(x=RFile('foo'), spam=RFile('bar'))),

1633

1682

` ('bar -x foo', NS(x=RFile('foo'), spam=RFile('bar'))),

1634

('-x - -', NS(x=eq_stdin, spam=eq_stdin)),

1683

('-x - -', NS(x=eq_bstdin, spam=eq_bstdin)),

1635

1684

` ]

1636

1685

1637

1686

`@@ -1658,8 +1707,9 @@ class TestFileTypeW(TempDirMixin, ParserTestCase):

1658

1707

`"""Test the FileType option/argument type for writing files"""

1659

1708

1660

1709

`def setUp(self):

1661

super(TestFileTypeW, self).setUp()

1710

super().setUp()

1662

1711

`self.create_readonly_file('readonly')

1712

self.create_writable_file('writable')

1663

1713

1664

1714

`argument_signatures = [

1665

1715

`Sig('-x', type=argparse.FileType('w')),

`@@ -1668,13 +1718,37 @@ def setUp(self):

1668

1718

`failures = ['-x', '', 'readonly']

1669

1719

`successes = [

1670

1720

` ('foo', NS(x=None, spam=WFile('foo'))),

1721

('writable', NS(x=None, spam=WFile('writable'))),

1671

1722

` ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),

1672

1723

` ('bar -x foo', NS(x=WFile('foo'), spam=WFile('bar'))),

1673

1724

` ('-x - -', NS(x=eq_stdout, spam=eq_stdout)),

1674

1725

` ]

1675

1726

1727

@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,

1728

"non-root user required")

1729

class TestFileTypeX(TempDirMixin, ParserTestCase):

1730

"""Test the FileType option/argument type for writing new files only"""

1731

+

1732

def setUp(self):

1733

super().setUp()

1734

self.create_readonly_file('readonly')

1735

self.create_writable_file('writable')

1736

+

1737

argument_signatures = [

1738

Sig('-x', type=argparse.FileType('x')),

1739

Sig('spam', type=argparse.FileType('x')),

1740

]

1741

failures = ['-x', '', 'readonly', 'writable']

1742

successes = [

1743

('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),

1744

('-x - -', NS(x=eq_stdout, spam=eq_stdout)),

1745

]

1746

+

1676

1747

1748

@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,

1749

"non-root user required")

1677

1750

`class TestFileTypeWB(TempDirMixin, ParserTestCase):

1751

"""Test the FileType option/argument type for writing binary files"""

1678

1752

1679

1753

`argument_signatures = [

1680

1754

`Sig('-x', type=argparse.FileType('wb')),

`@@ -1685,7 +1759,22 @@ class TestFileTypeWB(TempDirMixin, ParserTestCase):

1685

1759

` ('foo', NS(x=None, spam=WFile('foo'))),

1686

1760

` ('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),

1687

1761

` ('bar -x foo', NS(x=WFile('foo'), spam=WFile('bar'))),

1688

('-x - -', NS(x=eq_stdout, spam=eq_stdout)),

1762

('-x - -', NS(x=eq_bstdout, spam=eq_bstdout)),

1763

]

1764

+

1765

+

1766

@unittest.skipIf(hasattr(os, 'geteuid') and os.geteuid() == 0,

1767

"non-root user required")

1768

class TestFileTypeXB(TestFileTypeX):

1769

"Test the FileType option/argument type for writing new binary files only"

1770

+

1771

argument_signatures = [

1772

Sig('-x', type=argparse.FileType('xb')),

1773

Sig('spam', type=argparse.FileType('xb')),

1774

]

1775

successes = [

1776

('-x foo bar', NS(x=WFile('foo'), spam=WFile('bar'))),

1777

('-x - -', NS(x=eq_bstdout, spam=eq_bstdout)),

1689

1778

` ]

1690

1779

1691

1780