1"""Test correct operation of the print function.
2"""
3
4# In 2.6, this gives us the behavior we want.  In 3.0, it has
5#  no function, but it still must parse correctly.
6from __future__ import print_function
7
8import unittest
9from test import test_support
10
11from StringIO import StringIO
12
13NotDefined = object()
14
15# A dispatch table all 8 combinations of providing
16#  sep, end, and file
17# I use this machinery so that I'm not just passing default
18#  values to print, I'm either passing or not passing in the
19#  arguments
20dispatch = {
21    (False, False, False):
22     lambda args, sep, end, file: print(*args),
23    (False, False, True):
24     lambda args, sep, end, file: print(file=file, *args),
25    (False, True,  False):
26     lambda args, sep, end, file: print(end=end, *args),
27    (False, True,  True):
28     lambda args, sep, end, file: print(end=end, file=file, *args),
29    (True,  False, False):
30     lambda args, sep, end, file: print(sep=sep, *args),
31    (True,  False, True):
32     lambda args, sep, end, file: print(sep=sep, file=file, *args),
33    (True,  True,  False):
34     lambda args, sep, end, file: print(sep=sep, end=end, *args),
35    (True,  True,  True):
36     lambda args, sep, end, file: print(sep=sep, end=end, file=file, *args),
37    }
38
39# Class used to test __str__ and print
40class ClassWith__str__:
41    def __init__(self, x):
42        self.x = x
43    def __str__(self):
44        return self.x
45
46class TestPrint(unittest.TestCase):
47    def check(self, expected, args,
48            sep=NotDefined, end=NotDefined, file=NotDefined):
49        # Capture sys.stdout in a StringIO.  Call print with args,
50        #  and with sep, end, and file, if they're defined.  Result
51        #  must match expected.
52
53        # Look up the actual function to call, based on if sep, end, and file
54        #  are defined
55        fn = dispatch[(sep is not NotDefined,
56                       end is not NotDefined,
57                       file is not NotDefined)]
58
59        with test_support.captured_stdout() as t:
60            fn(args, sep, end, file)
61
62        self.assertEqual(t.getvalue(), expected)
63
64    def test_print(self):
65        def x(expected, args, sep=NotDefined, end=NotDefined):
66            # Run the test 2 ways: not using file, and using
67            #  file directed to a StringIO
68
69            self.check(expected, args, sep=sep, end=end)
70
71            # When writing to a file, stdout is expected to be empty
72            o = StringIO()
73            self.check('', args, sep=sep, end=end, file=o)
74
75            # And o will contain the expected output
76            self.assertEqual(o.getvalue(), expected)
77
78        x('\n', ())
79        x('a\n', ('a',))
80        x('None\n', (None,))
81        x('1 2\n', (1, 2))
82        x('1   2\n', (1, ' ', 2))
83        x('1*2\n', (1, 2), sep='*')
84        x('1 s', (1, 's'), end='')
85        x('a\nb\n', ('a', 'b'), sep='\n')
86        x('1.01', (1.0, 1), sep='', end='')
87        x('1*a*1.3+', (1, 'a', 1.3), sep='*', end='+')
88        x('a\n\nb\n', ('a\n', 'b'), sep='\n')
89        x('\0+ +\0\n', ('\0', ' ', '\0'), sep='+')
90
91        x('a\n b\n', ('a\n', 'b'))
92        x('a\n b\n', ('a\n', 'b'), sep=None)
93        x('a\n b\n', ('a\n', 'b'), end=None)
94        x('a\n b\n', ('a\n', 'b'), sep=None, end=None)
95
96        x('*\n', (ClassWith__str__('*'),))
97        x('abc 1\n', (ClassWith__str__('abc'), 1))
98
99        # 2.x unicode tests
100        x(u'1 2\n', ('1', u'2'))
101        x(u'u\1234\n', (u'u\1234',))
102        x(u'  abc 1\n', (' ', ClassWith__str__(u'abc'), 1))
103
104        # errors
105        self.assertRaises(TypeError, print, '', sep=3)
106        self.assertRaises(TypeError, print, '', end=3)
107        self.assertRaises(AttributeError, print, '', file='')
108
109    def test_mixed_args(self):
110        # If a unicode arg is passed, sep and end should be unicode, too.
111        class Recorder(object):
112
113            def __init__(self, must_be_unicode):
114                self.buf = []
115                self.force_unicode = must_be_unicode
116
117            def write(self, what):
118                if self.force_unicode and not isinstance(what, unicode):
119                    raise AssertionError("{0!r} is not unicode".format(what))
120                self.buf.append(what)
121
122        buf = Recorder(True)
123        print(u'hi', file=buf)
124        self.assertEqual(u''.join(buf.buf), 'hi\n')
125        del buf.buf[:]
126        print(u'hi', u'nothing', file=buf)
127        self.assertEqual(u''.join(buf.buf), 'hi nothing\n')
128        buf = Recorder(False)
129        print('hi', 'bye', end=u'\n', file=buf)
130        self.assertIsInstance(buf.buf[1], unicode)
131        self.assertIsInstance(buf.buf[3], unicode)
132        del buf.buf[:]
133        print(sep=u'x', file=buf)
134        self.assertIsInstance(buf.buf[-1], unicode)
135
136
137def test_main():
138    test_support.run_unittest(TestPrint)
139
140if __name__ == "__main__":
141    test_main()
142