1import dis
2import sys
3import textwrap
4import unittest
5
6from test.support import os_helper, verbose
7from test.support.script_helper import assert_python_ok
8
9def example():
10    x = []
11    for i in range(1):
12        x.append(i)
13    x = "this is"
14    y = "an example"
15    print(x, y)
16
17Py_DEBUG = hasattr(sys, 'gettotalrefcount')
18
19@unittest.skipUnless(Py_DEBUG, "lltrace requires Py_DEBUG")
20class TestLLTrace(unittest.TestCase):
21
22    def run_code(self, code):
23        code = textwrap.dedent(code).strip()
24        with open(os_helper.TESTFN, 'w', encoding='utf-8') as fd:
25            self.addCleanup(os_helper.unlink, os_helper.TESTFN)
26            fd.write(code)
27        status, stdout, stderr = assert_python_ok(os_helper.TESTFN)
28        self.assertEqual(stderr, b"")
29        self.assertEqual(status, 0)
30        result = stdout.decode('utf-8')
31        if verbose:
32            print("\n\n--- code ---")
33            print(code)
34            print("\n--- stdout ---")
35            print(result)
36            print()
37        return result
38
39    def test_lltrace(self):
40        stdout = self.run_code("""
41            def dont_trace_1():
42                a = "a"
43                a = 10 * a
44            def trace_me():
45                for i in range(3):
46                    +i
47            def dont_trace_2():
48                x = 42
49                y = -x
50            dont_trace_1()
51            __lltrace__ = 1
52            trace_me()
53            del __lltrace__
54            dont_trace_2()
55        """)
56        self.assertIn("GET_ITER", stdout)
57        self.assertIn("FOR_ITER", stdout)
58        self.assertIn("UNARY_POSITIVE", stdout)
59        self.assertIn("POP_TOP", stdout)
60        self.assertNotIn("BINARY_OP", stdout)
61        self.assertNotIn("UNARY_NEGATIVE", stdout)
62
63        self.assertIn("'trace_me' in module '__main__'", stdout)
64        self.assertNotIn("dont_trace_1", stdout)
65        self.assertNotIn("'dont_trace_2' in module", stdout)
66
67    def test_lltrace_different_module(self):
68        stdout = self.run_code("""
69            from test import test_lltrace
70            test_lltrace.__lltrace__ = 1
71            test_lltrace.example()
72        """)
73        self.assertIn("'example' in module 'test.test_lltrace'", stdout)
74        self.assertIn('LOAD_CONST', stdout)
75        self.assertIn('FOR_ITER', stdout)
76        self.assertIn('this is an example', stdout)
77
78        # check that offsets match the output of dis.dis()
79        instr_map = {i.offset: i for i in dis.get_instructions(example)}
80        for line in stdout.splitlines():
81            offset, colon, opname_oparg = line.partition(":")
82            if not colon:
83                continue
84            offset = int(offset)
85            opname_oparg = opname_oparg.split()
86            if len(opname_oparg) == 2:
87                opname, oparg = opname_oparg
88                oparg = int(oparg)
89            else:
90                (opname,) = opname_oparg
91                oparg = None
92            self.assertEqual(instr_map[offset].opname, opname)
93            self.assertEqual(instr_map[offset].arg, oparg)
94
95    def test_lltrace_does_not_crash_on_subscript_operator(self):
96        # If this test fails, it will reproduce a crash reported as
97        # bpo-34113. The crash happened at the command line console of
98        # debug Python builds with __lltrace__ enabled (only possible in console),
99        # when the internal Python stack was negatively adjusted
100        stdout = self.run_code("""
101            import code
102
103            console = code.InteractiveConsole()
104            console.push('__lltrace__ = 1')
105            console.push('a = [1, 2, 3]')
106            console.push('a[0] = 1')
107            print('unreachable if bug exists')
108        """)
109        self.assertIn("unreachable if bug exists", stdout)
110
111if __name__ == "__main__":
112    unittest.main()
113