1 import unittest 2 from io import StringIO 3 4 from test import support 5 6 NotDefined = object() 7 8 # A dispatch table all 8 combinations of providing 9 # sep, end, and file. 10 # I use this machinery so that I'm not just passing default 11 # values to print, I'm either passing or not passing in the 12 # arguments. 13 dispatch = { 14 (False, False, False): 15 lambda args, sep, end, file: print(*args), 16 (False, False, True): 17 lambda args, sep, end, file: print(file=file, *args), 18 (False, True, False): 19 lambda args, sep, end, file: print(end=end, *args), 20 (False, True, True): 21 lambda args, sep, end, file: print(end=end, file=file, *args), 22 (True, False, False): 23 lambda args, sep, end, file: print(sep=sep, *args), 24 (True, False, True): 25 lambda args, sep, end, file: print(sep=sep, file=file, *args), 26 (True, True, False): 27 lambda args, sep, end, file: print(sep=sep, end=end, *args), 28 (True, True, True): 29 lambda args, sep, end, file: print(sep=sep, end=end, file=file, *args), 30 } 31 32 33 # Class used to test __str__ and print 34 class ClassWith__str__: 35 def __init__(self, x): 36 self.x = x 37 38 def __str__(self): 39 return self.x 40 41 42 class TestPrint(unittest.TestCase): 43 """Test correct operation of the print function.""" 44 45 def check(self, expected, args, 46 sep=NotDefined, end=NotDefined, file=NotDefined): 47 # Capture sys.stdout in a StringIO. Call print with args, 48 # and with sep, end, and file, if they're defined. Result 49 # must match expected. 50 51 # Look up the actual function to call, based on if sep, end, 52 # and file are defined. 53 fn = dispatch[(sep is not NotDefined, 54 end is not NotDefined, 55 file is not NotDefined)] 56 57 with support.captured_stdout() as t: 58 fn(args, sep, end, file) 59 60 self.assertEqual(t.getvalue(), expected) 61 62 def test_print(self): 63 def x(expected, args, sep=NotDefined, end=NotDefined): 64 # Run the test 2 ways: not using file, and using 65 # file directed to a StringIO. 66 67 self.check(expected, args, sep=sep, end=end) 68 69 # When writing to a file, stdout is expected to be empty 70 o = StringIO() 71 self.check('', args, sep=sep, end=end, file=o) 72 73 # And o will contain the expected output 74 self.assertEqual(o.getvalue(), expected) 75 76 x('\n', ()) 77 x('a\n', ('a',)) 78 x('None\n', (None,)) 79 x('1 2\n', (1, 2)) 80 x('1 2\n', (1, ' ', 2)) 81 x('1*2\n', (1, 2), sep='*') 82 x('1 s', (1, 's'), end='') 83 x('a\nb\n', ('a', 'b'), sep='\n') 84 x('1.01', (1.0, 1), sep='', end='') 85 x('1*a*1.3+', (1, 'a', 1.3), sep='*', end='+') 86 x('a\n\nb\n', ('a\n', 'b'), sep='\n') 87 x('\0+ +\0\n', ('\0', ' ', '\0'), sep='+') 88 89 x('a\n b\n', ('a\n', 'b')) 90 x('a\n b\n', ('a\n', 'b'), sep=None) 91 x('a\n b\n', ('a\n', 'b'), end=None) 92 x('a\n b\n', ('a\n', 'b'), sep=None, end=None) 93 94 x('*\n', (ClassWith__str__('*'),)) 95 x('abc 1\n', (ClassWith__str__('abc'), 1)) 96 97 # errors 98 self.assertRaises(TypeError, print, '', sep=3) 99 self.assertRaises(TypeError, print, '', end=3) 100 self.assertRaises(AttributeError, print, '', file='') 101 102 def test_print_flush(self): 103 # operation of the flush flag 104 class filelike: 105 def __init__(self): 106 self.written = '' 107 self.flushed = 0 108 109 def write(self, str): 110 self.written += str 111 112 def flush(self): 113 self.flushed += 1 114 115 f = filelike() 116 print(1, file=f, end='', flush=True) 117 print(2, file=f, end='', flush=True) 118 print(3, file=f, flush=False) 119 self.assertEqual(f.written, '123\n') 120 self.assertEqual(f.flushed, 2) 121 122 # ensure exceptions from flush are passed through 123 class noflush: 124 def write(self, str): 125 pass 126 127 def flush(self): 128 raise RuntimeError 129 self.assertRaises(RuntimeError, print, 1, file=noflush(), flush=True) 130 131 if __name__ == "__main__": 132 unittest.main() 133
