bpo-14156: Make argparse.FileType work correctly for binary file mode… · python/cpython@ee18df4 (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.support import os_helper

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', encoding="utf-8") 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 = {}

`@@ -1633,7 +1682,7 @@ def setUp(self):

1633

1682

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

1634

1683

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

1635

1684

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

1636

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

1685

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

1637

1686

` ]

1638

1687

1639

1688

`@@ -1660,8 +1709,9 @@ class TestFileTypeW(TempDirMixin, ParserTestCase):

1660

1709

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

1661

1710

1662

1711

`def setUp(self):

1663

super(TestFileTypeW, self).setUp()

1712

super().setUp()

1664

1713

`self.create_readonly_file('readonly')

1714

self.create_writable_file('writable')

1665

1715

1666

1716

`argument_signatures = [

1667

1717

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

`@@ -1670,13 +1720,37 @@ def setUp(self):

1670

1720

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

1671

1721

`successes = [

1672

1722

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

1723

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

1673

1724

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

1674

1725

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

1675

1726

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

1676

1727

` ]

1677

1728

1729

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

1730

"non-root user required")

1731

class TestFileTypeX(TempDirMixin, ParserTestCase):

1732

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

1733

+

1734

def setUp(self):

1735

super().setUp()

1736

self.create_readonly_file('readonly')

1737

self.create_writable_file('writable')

1738

+

1739

argument_signatures = [

1740

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

1741

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

1742

]

1743

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

1744

successes = [

1745

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

1746

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

1747

]

1748

+

1678

1749

1750

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

1751

"non-root user required")

1679

1752

`class TestFileTypeWB(TempDirMixin, ParserTestCase):

1753

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

1680

1754

1681

1755

`argument_signatures = [

1682

1756

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

`@@ -1687,7 +1761,22 @@ class TestFileTypeWB(TempDirMixin, ParserTestCase):

1687

1761

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

1688

1762

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

1689

1763

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

1690

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

1764

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

1765

]

1766

+

1767

+

1768

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

1769

"non-root user required")

1770

class TestFileTypeXB(TestFileTypeX):

1771

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

1772

+

1773

argument_signatures = [

1774

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

1775

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

1776

]

1777

successes = [

1778

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

1779

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

1691

1780

` ]

1692

1781

1693

1782