GH-129386: Add `test.support.reset_code` (GH-129486)

6 files changed, 22 insertions, 9 deletions

diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py
index 6436753..230bb24 100644
--- a/Lib/test/support/__init__.py
+++ b/Lib/test/support/__init__.py
@@ -66,6 +66,7 @@ __all__ = [
     "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
index d3aea37..02e534c 100644
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -9,7 +9,8 @@ import os
 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 @@ from _testinternalcapi import TIER2_THRESHOLD
 @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
index bba2ac8..e99289c 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -15,7 +15,7 @@ import types
 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 @@ class DisTests(DisTestBase):
             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
index 7222137..cd65496 100644
--- a/Lib/test/test_embed.py
+++ b/Lib/test/test_embed.py
@@ -391,6 +391,7 @@ class EmbeddingTests(EmbeddingTestsMixin, unittest.TestCase):
             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 @@ class EmbeddingTests(EmbeddingTestsMixin, unittest.TestCase):
             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
index 2defe74..87de4c9 100644
--- a/Lib/test/test_opcache.py
+++ b/Lib/test/test_opcache.py
@@ -6,7 +6,7 @@ import types
 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 @@ class TestRacesDoNotCrash(TestBase):
             # 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 @@ class TestSpecializer(TestBase):
         class C:
             pass
 
+        @reset_code
         def set_value(n):
             c = C()
             for i in range(n):
@@ -1577,6 +1578,7 @@ class TestSpecializer(TestBase):
         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
diff --git a/Misc/NEWS.d/next/Tests/2025-01-30-13-09-27.gh-issue-129386.iNtbEi.rst b/Misc/NEWS.d/next/Tests/2025-01-30-13-09-27.gh-issue-129386.iNtbEi.rst
new file mode 100644
index 0000000..a03f596
--- /dev/null
+++ b/Misc/NEWS.d/next/Tests/2025-01-30-13-09-27.gh-issue-129386.iNtbEi.rst
@@ -0,0 +1,2 @@
+Add ``test.support.reset_code``, which can be used to reset various
+bytecode-level optimizations and local instrumentation for a function.