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