@@ -2343,23 +2343,26 @@ def func():
23432343 names = ["func" , "outer" , "outer" , "inner" , "inner" , "outer" , "inner" ]
23442344 self .do_test (func , names )
23452345
2346- class TestOptimizerAPI (unittest .TestCase ):
23472346
2348- @contextlib .contextmanager
2349- def temporary_optimizer (self , opt ):
2350- _testinternalcapi .set_optimizer (opt )
2351- try :
2352- yield
2353- finally :
2354- _testinternalcapi .set_optimizer (None )
2347+ @contextlib .contextmanager
2348+ def temporary_optimizer (opt ):
2349+ _testinternalcapi .set_optimizer (opt )
2350+ try :
2351+ yield
2352+ finally :
2353+ _testinternalcapi .set_optimizer (None )
23552354
2356- @contextlib .contextmanager
2357- def clear_executors (self , func ):
2358- try :
2359- yield
2360- finally :
2361- #Clear executors
2362- func .__code__ = func .__code__ .replace ()
2355+ @contextlib .contextmanager
2356+ def clear_executors (func ):
2357+ # Clear executors in func before and after running a block
2358+ func .__code__ = func .__code__ .replace ()
2359+ try :
2360+ yield
2361+ finally :
2362+ func .__code__ = func .__code__ .replace ()
2363+
2364+
2365+ class TestOptimizerAPI (unittest .TestCase ):
23632366
23642367 def test_get_set_optimizer (self ):
23652368 self .assertEqual (_testinternalcapi .get_optimizer (), None )
@@ -2381,9 +2384,9 @@ def loop():
23812384
23822385 for repeat in range (5 ):
23832386 opt = _testinternalcapi .get_counter_optimizer ()
2384- with self . temporary_optimizer (opt ):
2387+ with temporary_optimizer (opt ):
23852388 self .assertEqual (opt .get_count (), 0 )
2386- with self . clear_executors (loop ):
2389+ with clear_executors (loop ):
23872390 loop ()
23882391 self .assertEqual (opt .get_count (), 1000 )
23892392
@@ -2409,7 +2412,7 @@ def long_loop():
24092412 long_loop = ns ['long_loop' ]
24102413
24112414 opt = _testinternalcapi .get_counter_optimizer ()
2412- with self . temporary_optimizer (opt ):
2415+ with temporary_optimizer (opt ):
24132416 self .assertEqual (opt .get_count (), 0 )
24142417 long_loop ()
24152418 self .assertEqual (opt .get_count (), 10 )
@@ -2418,24 +2421,27 @@ def long_loop():
24182421class TestUops (unittest .TestCase ):
24192422
24202423 def test_basic_loop (self ):
2421-
24222424 def testfunc (x ):
24232425 i = 0
24242426 while i < x :
24252427 i += 1
24262428
2427- testfunc (1000 )
2429+ opt = _testinternalcapi .get_uop_optimizer ()
2430+
2431+ with temporary_optimizer (opt ):
2432+ testfunc (1000 )
24282433
24292434 ex = None
2430- for offset in range (0 , 100 , 2 ):
2435+ for offset in range (0 , len ( testfunc . __code__ . co_code ) , 2 ):
24312436 try :
24322437 ex = _testinternalcapi .get_executor (testfunc .__code__ , offset )
24332438 break
24342439 except ValueError :
24352440 pass
2436- if ex is None :
2437- return
2438- self .assertIn ("SAVE_IP" , str (ex ))
2441+ self .assertIsNotNone (ex )
2442+ uops = {opname for opname , _ in ex }
2443+ self .assertIn ("SAVE_IP" , uops )
2444+ self .assertIn ("LOAD_FAST" , uops )
24392445
24402446
24412447if __name__ == "__main__" :
0 commit comments