gh-104584: Clean up and fix uops tests and fix crash

Lib/test/test_capi/test_misc.py

@@ -2343,23 +2343,26 @@ def func():
         names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
         self.do_test(func, names)
 
-class TestOptimizerAPI(unittest.TestCase):
 
-    @contextlib.contextmanager
-    def temporary_optimizer(self, opt):
-        _testinternalcapi.set_optimizer(opt)
-        try:
-            yield
-        finally:
-            _testinternalcapi.set_optimizer(None)
+@contextlib.contextmanager
+def temporary_optimizer(opt):
+    _testinternalcapi.set_optimizer(opt)
+    try:
+        yield
+    finally:
+        _testinternalcapi.set_optimizer(None)
 
-    @contextlib.contextmanager
-    def clear_executors(self, func):
-        try:
-            yield
-        finally:
-            #Clear executors
-            func.__code__ = func.__code__.replace()
+@contextlib.contextmanager
+def clear_executors(func):
+    # Clear executors in func before and after running a block
+    func.__code__ = func.__code__.replace()
+    try:
+        yield
+    finally:
+        func.__code__ = func.__code__.replace()
+
+
+class TestOptimizerAPI(unittest.TestCase):
 
     def test_get_set_optimizer(self):
         self.assertEqual(_testinternalcapi.get_optimizer(), None)
@@ -2381,9 +2384,9 @@ def loop():
 
         for repeat in range(5):
             opt = _testinternalcapi.get_counter_optimizer()
-            with self.temporary_optimizer(opt):
+            with temporary_optimizer(opt):
                 self.assertEqual(opt.get_count(), 0)
-                with self.clear_executors(loop):
+                with clear_executors(loop):
                     loop()
                 self.assertEqual(opt.get_count(), 1000)
 
@@ -2409,7 +2412,7 @@ def long_loop():
         long_loop = ns['long_loop']
 
         opt = _testinternalcapi.get_counter_optimizer()
-        with self.temporary_optimizer(opt):
+        with temporary_optimizer(opt):
             self.assertEqual(opt.get_count(), 0)
             long_loop()
             self.assertEqual(opt.get_count(), 10)
@@ -2418,24 +2421,27 @@ def long_loop():
 class TestUops(unittest.TestCase):
 
     def test_basic_loop(self):
-
         def testfunc(x):
             i = 0
             while i < x:
                 i += 1
 
-        testfunc(1000)
+        opt = _testinternalcapi.get_uop_optimizer()
+
+        with temporary_optimizer(opt):
+            testfunc(1000)
 
         ex = None
-        for offset in range(0, 100, 2):
+        for offset in range(0, len(testfunc.__code__.co_code), 2):
             try:
                 ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
                 break
             except ValueError:
                 pass
-        if ex is None:
-            return
-        self.assertIn("SAVE_IP", str(ex))
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _ in ex}
+        self.assertIn("SAVE_IP", uops)
+        self.assertIn("LOAD_FAST", uops)
 
 
 if __name__ == "__main__":

Python/optimizer.c

@@ -532,7 +532,7 @@ uop_optimize(
         return trace_length;
     }
     OBJECT_STAT_INC(optimization_traces_created);
-    _PyUOpExecutorObject *executor = (_PyUOpExecutorObject *)_PyObject_New(&UOpExecutor_Type);
+    _PyUOpExecutorObject *executor = PyObject_New(_PyUOpExecutorObject, &UOpExecutor_Type);
     if (executor == NULL) {
         return -1;
     }
@@ -542,18 +542,24 @@ uop_optimize(
     return 1;
 }
 
+static void
+uop_opt_dealloc(PyObject *self) {
+    PyObject_Free(self);
+}
+
 static PyTypeObject UOpOptimizer_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     .tp_name = "uop_optimizer",
     .tp_basicsize = sizeof(_PyOptimizerObject),
     .tp_itemsize = 0,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
+    .tp_dealloc = uop_opt_dealloc,
 };
 
 PyObject *
 PyUnstable_Optimizer_NewUOpOptimizer(void)
 {
-    _PyOptimizerObject *opt = (_PyOptimizerObject *)_PyObject_New(&UOpOptimizer_Type);
+    _PyOptimizerObject *opt = PyObject_New(_PyOptimizerObject, &UOpOptimizer_Type);
     if (opt == NULL) {
         return NULL;
     }