91 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4from mlir.ir import *5 6 7def run(f):8 print("\nTEST:", f.__name__)9 f()10 gc.collect()11 assert Context._get_live_count() == 012 13 14# CHECK-LABEL: TEST: testContextEnterExit15def testContextEnterExit():16 with Context() as ctx:17 assert Context.current is ctx18 assert Context.current is None19 20 21run(testContextEnterExit)22 23 24# CHECK-LABEL: TEST: testLocationEnterExit25def testLocationEnterExit():26 ctx1 = Context()27 with Location.unknown(ctx1) as loc1:28 assert Context.current is ctx129 assert Location.current is loc130 31 # Re-asserting the same context should not change the location.32 with ctx1:33 assert Context.current is ctx134 assert Location.current is loc135 # Asserting a different context should clear it.36 with Context() as ctx2:37 assert Context.current is ctx238 assert Location.current is None39 40 # And should restore.41 assert Context.current is ctx142 assert Location.current is loc143 44 # All should clear.45 assert Location.current is None46 47 48run(testLocationEnterExit)49 50 51# CHECK-LABEL: TEST: testInsertionPointEnterExit52def testInsertionPointEnterExit():53 ctx1 = Context()54 m = Module.create(Location.unknown(ctx1))55 ip = InsertionPoint(m.body)56 57 with ip:58 assert InsertionPoint.current is ip59 # Asserting a location from the same context should preserve.60 with Location.unknown(ctx1) as loc1:61 assert InsertionPoint.current is ip62 assert Location.current is loc163 # Location should clear.64 assert Location.current is None65 66 # Asserting the same Context should preserve.67 with ctx1:68 assert InsertionPoint.current is ip69 70 # Asserting a different context should clear it.71 with Context() as ctx2:72 assert Context.current is ctx273 try:74 _ = InsertionPoint.current75 except ValueError:76 pass77 else:78 assert False, "Expected exception"79 80 # All should clear.81 try:82 _ = InsertionPoint.current83 except ValueError as e:84 # CHECK: No current InsertionPoint85 print(e)86 else:87 assert False, "Expected exception"88 89 90run(testInsertionPointEnterExit)91