1import dis
2import sys
3from cStringIO import StringIO
4import unittest
5
6def disassemble(func):
7    f = StringIO()
8    tmp = sys.stdout
9    sys.stdout = f
10    dis.dis(func)
11    sys.stdout = tmp
12    result = f.getvalue()
13    f.close()
14    return result
15
16def dis_single(line):
17    return disassemble(compile(line, '', 'single'))
18
19class TestTranforms(unittest.TestCase):
20
21    def test_unot(self):
22        # UNARY_NOT POP_JUMP_IF_FALSE  -->  POP_JUMP_IF_TRUE
23        def unot(x):
24            if not x == 2:
25                del x
26        asm = disassemble(unot)
27        for elem in ('UNARY_NOT', 'POP_JUMP_IF_FALSE'):
28            self.assertNotIn(elem, asm)
29        self.assertIn('POP_JUMP_IF_TRUE', asm)
30
31    def test_elim_inversion_of_is_or_in(self):
32        for line, elem in (
33            ('not a is b', '(is not)',),
34            ('not a in b', '(not in)',),
35            ('not a is not b', '(is)',),
36            ('not a not in b', '(in)',),
37            ):
38            asm = dis_single(line)
39            self.assertIn(elem, asm)
40
41    def test_none_as_constant(self):
42        # LOAD_GLOBAL None  -->  LOAD_CONST None
43        def f(x):
44            None
45            return x
46        asm = disassemble(f)
47        for elem in ('LOAD_GLOBAL',):
48            self.assertNotIn(elem, asm)
49        for elem in ('LOAD_CONST', '(None)'):
50            self.assertIn(elem, asm)
51        def f():
52            'Adding a docstring made this test fail in Py2.5.0'
53            return None
54        self.assertIn('LOAD_CONST', disassemble(f))
55        self.assertNotIn('LOAD_GLOBAL', disassemble(f))
56
57    def test_while_one(self):
58        # Skip over:  LOAD_CONST trueconst  POP_JUMP_IF_FALSE xx
59        def f():
60            while 1:
61                pass
62            return list
63        asm = disassemble(f)
64        for elem in ('LOAD_CONST', 'POP_JUMP_IF_FALSE'):
65            self.assertNotIn(elem, asm)
66        for elem in ('JUMP_ABSOLUTE',):
67            self.assertIn(elem, asm)
68
69    def test_pack_unpack(self):
70        for line, elem in (
71            ('a, = a,', 'LOAD_CONST',),
72            ('a, b = a, b', 'ROT_TWO',),
73            ('a, b, c = a, b, c', 'ROT_THREE',),
74            ):
75            asm = dis_single(line)
76            self.assertIn(elem, asm)
77            self.assertNotIn('BUILD_TUPLE', asm)
78            self.assertNotIn('UNPACK_TUPLE', asm)
79
80    def test_folding_of_tuples_of_constants(self):
81        for line, elem in (
82            ('a = 1,2,3', '((1, 2, 3))'),
83            ('("a","b","c")', "(('a', 'b', 'c'))"),
84            ('a,b,c = 1,2,3', '((1, 2, 3))'),
85            ('(None, 1, None)', '((None, 1, None))'),
86            ('((1, 2), 3, 4)', '(((1, 2), 3, 4))'),
87            ):
88            asm = dis_single(line)
89            self.assertIn(elem, asm)
90            self.assertNotIn('BUILD_TUPLE', asm)
91
92        # Bug 1053819:  Tuple of constants misidentified when presented with:
93        # . . . opcode_with_arg 100   unary_opcode   BUILD_TUPLE 1  . . .
94        # The following would segfault upon compilation
95        def crater():
96            (~[
97                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
98                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
99                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
100                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
101                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
102                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
103                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
104                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
105                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
106                0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
107            ],)
108
109    def test_folding_of_binops_on_constants(self):
110        for line, elem in (
111            ('a = 2+3+4', '(9)'),                   # chained fold
112            ('"@"*4', "('@@@@')"),                  # check string ops
113            ('a="abc" + "def"', "('abcdef')"),      # check string ops
114            ('a = 3**4', '(81)'),                   # binary power
115            ('a = 3*4', '(12)'),                    # binary multiply
116            ('a = 13//4', '(3)'),                   # binary floor divide
117            ('a = 14%4', '(2)'),                    # binary modulo
118            ('a = 2+3', '(5)'),                     # binary add
119            ('a = 13-4', '(9)'),                    # binary subtract
120            ('a = (12,13)[1]', '(13)'),             # binary subscr
121            ('a = 13 << 2', '(52)'),                # binary lshift
122            ('a = 13 >> 2', '(3)'),                 # binary rshift
123            ('a = 13 & 7', '(5)'),                  # binary and
124            ('a = 13 ^ 7', '(10)'),                 # binary xor
125            ('a = 13 | 7', '(15)'),                 # binary or
126            ):
127            asm = dis_single(line)
128            self.assertIn(elem, asm, asm)
129            self.assertNotIn('BINARY_', asm)
130
131        # Verify that unfoldables are skipped
132        asm = dis_single('a=2+"b"')
133        self.assertIn('(2)', asm)
134        self.assertIn("('b')", asm)
135
136        # Verify that large sequences do not result from folding
137        asm = dis_single('a="x"*1000')
138        self.assertIn('(1000)', asm)
139
140    def test_binary_subscr_on_unicode(self):
141        # unicode strings don't get optimized
142        asm = dis_single('u"foo"[0]')
143        self.assertNotIn("(u'f')", asm)
144        self.assertIn('BINARY_SUBSCR', asm)
145        asm = dis_single('u"\u0061\uffff"[1]')
146        self.assertNotIn("(u'\\uffff')", asm)
147        self.assertIn('BINARY_SUBSCR', asm)
148
149        # out of range
150        asm = dis_single('u"fuu"[10]')
151        self.assertIn('BINARY_SUBSCR', asm)
152        # non-BMP char (see #5057)
153        asm = dis_single('u"\U00012345"[0]')
154        self.assertIn('BINARY_SUBSCR', asm)
155        asm = dis_single('u"\U00012345abcdef"[3]')
156        self.assertIn('BINARY_SUBSCR', asm)
157
158
159    def test_folding_of_unaryops_on_constants(self):
160        for line, elem in (
161            ('`1`', "('1')"),                       # unary convert
162            ('-0.5', '(-0.5)'),                     # unary negative
163            ('~-2', '(1)'),                         # unary invert
164        ):
165            asm = dis_single(line)
166            self.assertIn(elem, asm, asm)
167            self.assertNotIn('UNARY_', asm)
168
169        # Verify that unfoldables are skipped
170        for line, elem in (
171            ('-"abc"', "('abc')"),                  # unary negative
172            ('~"abc"', "('abc')"),                  # unary invert
173        ):
174            asm = dis_single(line)
175            self.assertIn(elem, asm, asm)
176            self.assertIn('UNARY_', asm)
177
178    def test_elim_extra_return(self):
179        # RETURN LOAD_CONST None RETURN  -->  RETURN
180        def f(x):
181            return x
182        asm = disassemble(f)
183        self.assertNotIn('LOAD_CONST', asm)
184        self.assertNotIn('(None)', asm)
185        self.assertEqual(asm.split().count('RETURN_VALUE'), 1)
186
187    def test_elim_jump_to_return(self):
188        # JUMP_FORWARD to RETURN -->  RETURN
189        def f(cond, true_value, false_value):
190            return true_value if cond else false_value
191        asm = disassemble(f)
192        self.assertNotIn('JUMP_FORWARD', asm)
193        self.assertNotIn('JUMP_ABSOLUTE', asm)
194        self.assertEqual(asm.split().count('RETURN_VALUE'), 2)
195
196    def test_elim_jump_after_return1(self):
197        # Eliminate dead code: jumps immediately after returns can't be reached
198        def f(cond1, cond2):
199            if cond1: return 1
200            if cond2: return 2
201            while 1:
202                return 3
203            while 1:
204                if cond1: return 4
205                return 5
206            return 6
207        asm = disassemble(f)
208        self.assertNotIn('JUMP_FORWARD', asm)
209        self.assertNotIn('JUMP_ABSOLUTE', asm)
210        self.assertEqual(asm.split().count('RETURN_VALUE'), 6)
211
212    def test_elim_jump_after_return2(self):
213        # Eliminate dead code: jumps immediately after returns can't be reached
214        def f(cond1, cond2):
215            while 1:
216                if cond1: return 4
217        asm = disassemble(f)
218        self.assertNotIn('JUMP_FORWARD', asm)
219        # There should be one jump for the while loop.
220        self.assertEqual(asm.split().count('JUMP_ABSOLUTE'), 1)
221        self.assertEqual(asm.split().count('RETURN_VALUE'), 2)
222
223
224def test_main(verbose=None):
225    import sys
226    from test import test_support
227    test_classes = (TestTranforms,)
228
229    with test_support.check_py3k_warnings(
230            ("backquote not supported", SyntaxWarning)):
231        test_support.run_unittest(*test_classes)
232
233        # verify reference counting
234        if verbose and hasattr(sys, "gettotalrefcount"):
235            import gc
236            counts = [None] * 5
237            for i in xrange(len(counts)):
238                test_support.run_unittest(*test_classes)
239                gc.collect()
240                counts[i] = sys.gettotalrefcount()
241            print counts
242
243if __name__ == "__main__":
244    test_main(verbose=True)
245