GH-104584: Assorted fixes for the optimizer API. (GH-105683)

* Add test for long loops

* Clear ENTER_EXECUTOR when deopting code objects.

1 files changed, 57 insertions, 8 deletions

diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
index ccec27d..f2aa2a0 100644
--- a/Lib/test/test_capi/test_misc.py
+++ b/Lib/test/test_capi/test_misc.py
@@ -2347,17 +2347,66 @@ class Test_Pep523API(unittest.TestCase):
 
 class TestOptimizerAPI(unittest.TestCase):
 
-    def test_counter_optimizer(self):
-        opt = _testinternalcapi.get_counter_optimizer()
-        self.assertEqual(opt.get_count(), 0)
+    @contextlib.contextmanager
+    def temporary_optimizer(self, opt):
+        _testinternalcapi.set_optimizer(opt)
         try:
-            _testinternalcapi.set_optimizer(opt)
-            self.assertEqual(opt.get_count(), 0)
-            for _ in range(1000):
-                pass
-            self.assertEqual(opt.get_count(), 1000)
+            yield
         finally:
             _testinternalcapi.set_optimizer(None)
 
+    @contextlib.contextmanager
+    def clear_executors(self, func):
+        try:
+            yield
+        finally:
+            #Clear executors
+            func.__code__ = func.__code__.replace()
+
+    def test_get_set_optimizer(self):
+        self.assertEqual(_testinternalcapi.get_optimizer(), None)
+        opt = _testinternalcapi.get_counter_optimizer()
+        _testinternalcapi.set_optimizer(opt)
+        self.assertEqual(_testinternalcapi.get_optimizer(), opt)
+        _testinternalcapi.set_optimizer(None)
+        self.assertEqual(_testinternalcapi.get_optimizer(), None)
+
+    def test_counter_optimizer(self):
+
+        def loop():
+            for _ in range(1000):
+                pass
+
+        for repeat in range(5):
+            opt = _testinternalcapi.get_counter_optimizer()
+            with self.temporary_optimizer(opt):
+                self.assertEqual(opt.get_count(), 0)
+                with self.clear_executors(loop):
+                    loop()
+                self.assertEqual(opt.get_count(), 1000)
+
+    def test_long_loop(self):
+        "Check that we aren't confused by EXTENDED_ARG"
+
+        def nop():
+            pass
+
+        def long_loop():
+            for _ in range(10):
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+
+        opt = _testinternalcapi.get_counter_optimizer()
+        with self.temporary_optimizer(opt):
+            self.assertEqual(opt.get_count(), 0)
+            long_loop()
+            self.assertEqual(opt.get_count(), 10)
+
+
 if __name__ == "__main__":
     unittest.main()