1import gc
2import sys
3import types
4import unittest
5import weakref
6
7from test import support
8
9
10class ClearTest(unittest.TestCase):
11    """
12    Tests for frame.clear().
13    """
14
15    def inner(self, x=5, **kwargs):
16        1/0
17
18    def outer(self, **kwargs):
19        try:
20            self.inner(**kwargs)
21        except ZeroDivisionError as e:
22            exc = e
23        return exc
24
25    def clear_traceback_frames(self, tb):
26        """
27        Clear all frames in a traceback.
28        """
29        while tb is not None:
30            tb.tb_frame.clear()
31            tb = tb.tb_next
32
33    def test_clear_locals(self):
34        class C:
35            pass
36        c = C()
37        wr = weakref.ref(c)
38        exc = self.outer(c=c)
39        del c
40        support.gc_collect()
41        # A reference to c is held through the frames
42        self.assertIsNot(None, wr())
43        self.clear_traceback_frames(exc.__traceback__)
44        support.gc_collect()
45        # The reference was released by .clear()
46        self.assertIs(None, wr())
47
48    def test_clear_generator(self):
49        endly = False
50        def g():
51            nonlocal endly
52            try:
53                yield
54                inner()
55            finally:
56                endly = True
57        gen = g()
58        next(gen)
59        self.assertFalse(endly)
60        # Clearing the frame closes the generator
61        gen.gi_frame.clear()
62        self.assertTrue(endly)
63
64    def test_clear_executing(self):
65        # Attempting to clear an executing frame is forbidden.
66        try:
67            1/0
68        except ZeroDivisionError as e:
69            f = e.__traceback__.tb_frame
70        with self.assertRaises(RuntimeError):
71            f.clear()
72        with self.assertRaises(RuntimeError):
73            f.f_back.clear()
74
75    def test_clear_executing_generator(self):
76        # Attempting to clear an executing generator frame is forbidden.
77        endly = False
78        def g():
79            nonlocal endly
80            try:
81                1/0
82            except ZeroDivisionError as e:
83                f = e.__traceback__.tb_frame
84                with self.assertRaises(RuntimeError):
85                    f.clear()
86                with self.assertRaises(RuntimeError):
87                    f.f_back.clear()
88                yield f
89            finally:
90                endly = True
91        gen = g()
92        f = next(gen)
93        self.assertFalse(endly)
94        # Clearing the frame closes the generator
95        f.clear()
96        self.assertTrue(endly)
97
98    @support.cpython_only
99    def test_clear_refcycles(self):
100        # .clear() doesn't leave any refcycle behind
101        with support.disable_gc():
102            class C:
103                pass
104            c = C()
105            wr = weakref.ref(c)
106            exc = self.outer(c=c)
107            del c
108            self.assertIsNot(None, wr())
109            self.clear_traceback_frames(exc.__traceback__)
110            self.assertIs(None, wr())
111
112
113class FrameLocalsTest(unittest.TestCase):
114    """
115    Tests for the .f_locals attribute.
116    """
117
118    def make_frames(self):
119        def outer():
120            x = 5
121            y = 6
122            def inner():
123                z = x + 2
124                1/0
125                t = 9
126            return inner()
127        try:
128            outer()
129        except ZeroDivisionError as e:
130            tb = e.__traceback__
131            frames = []
132            while tb:
133                frames.append(tb.tb_frame)
134                tb = tb.tb_next
135        return frames
136
137    def test_locals(self):
138        f, outer, inner = self.make_frames()
139        outer_locals = outer.f_locals
140        self.assertIsInstance(outer_locals.pop('inner'), types.FunctionType)
141        self.assertEqual(outer_locals, {'x': 5, 'y': 6})
142        inner_locals = inner.f_locals
143        self.assertEqual(inner_locals, {'x': 5, 'z': 7})
144
145    def test_clear_locals(self):
146        # Test f_locals after clear() (issue #21897)
147        f, outer, inner = self.make_frames()
148        outer.clear()
149        inner.clear()
150        self.assertEqual(outer.f_locals, {})
151        self.assertEqual(inner.f_locals, {})
152
153    def test_locals_clear_locals(self):
154        # Test f_locals before and after clear() (to exercise caching)
155        f, outer, inner = self.make_frames()
156        outer.f_locals
157        inner.f_locals
158        outer.clear()
159        inner.clear()
160        self.assertEqual(outer.f_locals, {})
161        self.assertEqual(inner.f_locals, {})
162
163
164if __name__ == "__main__":
165    unittest.main()
166