Add test.support.reset_code

python · encukou · Jan 31, 2025 · Jan 30, 2025 · Jan 30, 2025 · Jan 30, 2025
commit 313708b126a68886a0c904e90062032914f11c73
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
@@ -66,6 +66,7 @@
     "BrokenIter",
     "in_systemd_nspawn_sync_suppressed",
     "run_no_yield_async_fn", "run_yielding_async_fn", "async_yield",
+    "reset_code",
     ]
 
 
@@ -1286,6 +1287,12 @@ def requires_specialization_ft(test):
         _opcode.ENABLE_SPECIALIZATION_FT, "requires specialization")(test)
 
 
+def reset_code(f: types.FunctionType) -> types.FunctionType:
+    """Clear all specializations, local instrumentation, and JIT code for the given function."""
+    f.__code__ = f.__code__.replace()
+    return f
+
+
 #=======================================================================
 # Check for the presence of docstrings.
 

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
@@ -9,7 +9,8 @@
 import _opcode
 
 from test.support import (script_helper, requires_specialization,
-                          import_helper, Py_GIL_DISABLED, requires_jit_enabled)
+                          import_helper, Py_GIL_DISABLED, requires_jit_enabled,
+                          reset_code)
 
 _testinternalcapi = import_helper.import_module("_testinternalcapi")
 
@@ -19,11 +20,11 @@
 @contextlib.contextmanager
 def clear_executors(func):
     # Clear executors in func before and after running a block
-    func.__code__ = func.__code__.replace()
+    reset_code(func)
     try:
         yield
     finally:
-        func.__code__ = func.__code__.replace()
+        reset_code(func)
 
 
 def get_first_executor(func):

diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
@@ -15,7 +15,7 @@
 import unittest
 from test.support import (captured_stdout, requires_debug_ranges,
                           requires_specialization, cpython_only,
-                          os_helper, import_helper)
+                          os_helper, import_helper, reset_code)
 from test.support.bytecode_helper import BytecodeTestCase
 
 
@@ -1356,7 +1356,7 @@ def f():
             self.code_quicken(f)
         else:
             # "copy" the code to un-quicken it:
-            f.__code__ = f.__code__.replace()
+            reset_code(f)
         for instruction in _unroll_caches_as_Instructions(dis.get_instructions(
             f, show_caches=True, adaptive=adaptive
         ), show_caches=True):

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
@@ -391,6 +391,7 @@ def test_specialized_static_code_gets_unspecialized_at_Py_FINALIZE(self):
             import importlib._bootstrap
             import opcode
             import test.test_dis
+            import test.support
 
             def is_specialized(f):
                 for instruction in dis.get_instructions(f, adaptive=True):
@@ -409,7 +410,7 @@ def is_specialized(f):
             func = importlib._bootstrap._handle_fromlist
 
             # "copy" the code to un-specialize it:
-            func.__code__ = func.__code__.replace()
+            test.support.reset_code(func)
 
             assert not is_specialized(func), "specialized instructions found"
 

diff --git a/Lib/test/test_opcache.py b/Lib/test/test_opcache.py
@@ -6,7 +6,7 @@
 import unittest
 from test.support import (threading_helper, check_impl_detail,
                           requires_specialization, requires_specialization_ft,
-                          cpython_only, requires_jit_disabled)
+                          cpython_only, requires_jit_disabled, reset_code)
 from test.support.import_helper import import_module
 
 # Skip this module on other interpreters, it is cpython specific:
@@ -579,9 +579,9 @@ def assert_races_do_not_crash(
             # Reset:
             if check_items:
                 for item in items:
-                    item.__code__ = item.__code__.replace()
+                    reset_code(item)
             else:
-                read.__code__ = read.__code__.replace()
+                reset_code(read)
             # Specialize:
             for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
                 read(items)
@@ -1552,6 +1552,7 @@ def test_store_attr_instance_value(self):
         class C:
             pass
 
+        @reset_code
         def set_value(n):
             c = C()
             for i in range(n):
@@ -1577,6 +1578,7 @@ class C:
         for i in range(_testinternalcapi.SHARED_KEYS_MAX_SIZE - 1):
             setattr(c, f"_{i}", None)
 
+        @reset_code
         def set_value(n):
             for i in range(n):
                 c.x = i