@@ -2343,23 +2343,26 @@ def func():
2343
2343
names = ["func" , "outer" , "outer" , "inner" , "inner" , "outer" , "inner" ]
2344
2344
self .do_test (func , names )
2345
2345
2346
- class TestOptimizerAPI (unittest .TestCase ):
2347
2346
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 )
2355
2354
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 ):
2363
2366
2364
2367
def test_get_set_optimizer (self ):
2365
2368
self .assertEqual (_testinternalcapi .get_optimizer (), None )
@@ -2381,9 +2384,9 @@ def loop():
2381
2384
2382
2385
for repeat in range (5 ):
2383
2386
opt = _testinternalcapi .get_counter_optimizer ()
2384
- with self . temporary_optimizer (opt ):
2387
+ with temporary_optimizer (opt ):
2385
2388
self .assertEqual (opt .get_count (), 0 )
2386
- with self . clear_executors (loop ):
2389
+ with clear_executors (loop ):
2387
2390
loop ()
2388
2391
self .assertEqual (opt .get_count (), 1000 )
2389
2392
@@ -2409,7 +2412,7 @@ def long_loop():
2409
2412
long_loop = ns ['long_loop' ]
2410
2413
2411
2414
opt = _testinternalcapi .get_counter_optimizer ()
2412
- with self . temporary_optimizer (opt ):
2415
+ with temporary_optimizer (opt ):
2413
2416
self .assertEqual (opt .get_count (), 0 )
2414
2417
long_loop ()
2415
2418
self .assertEqual (opt .get_count (), 10 )
@@ -2418,24 +2421,27 @@ def long_loop():
2418
2421
class TestUops (unittest .TestCase ):
2419
2422
2420
2423
def test_basic_loop (self ):
2421
-
2422
2424
def testfunc (x ):
2423
2425
i = 0
2424
2426
while i < x :
2425
2427
i += 1
2426
2428
2427
- testfunc (1000 )
2429
+ opt = _testinternalcapi .get_uop_optimizer ()
2430
+
2431
+ with temporary_optimizer (opt ):
2432
+ testfunc (1000 )
2428
2433
2429
2434
ex = None
2430
- for offset in range (0 , 100 , 2 ):
2435
+ for offset in range (0 , len ( testfunc . __code__ . co_code ) , 2 ):
2431
2436
try :
2432
2437
ex = _testinternalcapi .get_executor (testfunc .__code__ , offset )
2433
2438
break
2434
2439
except ValueError :
2435
2440
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 )
2439
2445
2440
2446
2441
2447
if __name__ == "__main__" :
0 commit comments