xref: /llvm-project/mlir/test/python/ir/context_managers.py (revision 225648e91ccd951eab9a4ab3200248d5617df1cc) 
1# RUN: %PYTHON %s | FileCheck %s
2
3import gc
4from mlir.ir import *
5
6
7def run(f):
8    print("\nTEST:", f.__name__)
9    f()
10    gc.collect()
11    assert Context._get_live_count() == 0
12
13
14# CHECK-LABEL: TEST: testContextEnterExit
15def testContextEnterExit():
16    with Context() as ctx:
17        assert Context.current is ctx
18    assert Context.current is None
19
20
21run(testContextEnterExit)
22
23
24# CHECK-LABEL: TEST: testLocationEnterExit
25def testLocationEnterExit():
26    ctx1 = Context()
27    with Location.unknown(ctx1) as loc1:
28        assert Context.current is ctx1
29        assert Location.current is loc1
30
31        # Re-asserting the same context should not change the location.
32        with ctx1:
33            assert Context.current is ctx1
34            assert Location.current is loc1
35            # Asserting a different context should clear it.
36            with Context() as ctx2:
37                assert Context.current is ctx2
38                try:
39                    _ = Location.current
40                except ValueError:
41                    pass
42                else:
43                    assert False, "Expected exception"
44
45            # And should restore.
46            assert Context.current is ctx1
47            assert Location.current is loc1
48
49    # All should clear.
50    try:
51        _ = Location.current
52    except ValueError as e:
53        # CHECK: No current Location
54        print(e)
55    else:
56        assert False, "Expected exception"
57
58
59run(testLocationEnterExit)
60
61
62# CHECK-LABEL: TEST: testInsertionPointEnterExit
63def testInsertionPointEnterExit():
64    ctx1 = Context()
65    m = Module.create(Location.unknown(ctx1))
66    ip = InsertionPoint(m.body)
67
68    with ip:
69        assert InsertionPoint.current is ip
70        # Asserting a location from the same context should preserve.
71        with Location.unknown(ctx1) as loc1:
72            assert InsertionPoint.current is ip
73            assert Location.current is loc1
74        # Location should clear.
75        try:
76            _ = Location.current
77        except ValueError:
78            pass
79        else:
80            assert False, "Expected exception"
81
82        # Asserting the same Context should preserve.
83        with ctx1:
84            assert InsertionPoint.current is ip
85
86        # Asserting a different context should clear it.
87        with Context() as ctx2:
88            assert Context.current is ctx2
89            try:
90                _ = InsertionPoint.current
91            except ValueError:
92                pass
93            else:
94                assert False, "Expected exception"
95
96    # All should clear.
97    try:
98        _ = InsertionPoint.current
99    except ValueError as e:
100        # CHECK: No current InsertionPoint
101        print(e)
102    else:
103        assert False, "Expected exception"
104
105
106run(testInsertionPointEnterExit)
107