python · markshannon · May 4, 2022 · May 3, 2022 · May 3, 2022 · May 3, 2022
diff --git a/Lib/test/test_capi.py b/Lib/test/test_capi.py
@@ -1142,5 +1142,33 @@ def test_frame_get_generator(self):
         self.assertIs(gen, _testcapi.frame_getgenerator(frame))
 
 
+SUFFICIENT_TO_SPECIALIZE = 50
+
+
+class Test_Pep523API(unittest.TestCase):
+
+    def do_test(self, func):
+        calls = []
+        _testinternalcapi.set_eval_frame_record(calls)
+        for _ in range(SUFFICIENT_TO_SPECIALIZE):
+            func()
+        _testinternalcapi.set_eval_frame_default()
+        self.assertEqual(len(calls), SUFFICIENT_TO_SPECIALIZE)
+        for name in calls:
+            self.assertEqual(name, func.__name__)
+
+    def test_intercept_before_specialize(self):
+        def func1():
+            pass
+        self.do_test(func1)
+
+    def test_specialize_before_intercept(self):
+        def func2():
+            pass
+        for _ in range(SUFFICIENT_TO_SPECIALIZE):
+            func2()
+        self.do_test(func2)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -15,14 +15,15 @@
 #include "pycore_atomic_funcs.h" // _Py_atomic_int_get()
 #include "pycore_bitutils.h"     // _Py_bswap32()
 #include "pycore_fileutils.h"    // _Py_normpath
+#include "pycore_frame.h"        // _PyInterpreterFrame
 #include "pycore_gc.h"           // PyGC_Head
 #include "pycore_hashtable.h"    // _Py_hashtable_new()
 #include "pycore_initconfig.h"   // _Py_GetConfigsAsDict()
 #include "pycore_pathconfig.h"   // _PyPathConfig_ClearGlobal()
 #include "pycore_interp.h"       // _PyInterpreterState_GetConfigCopy()
 #include "pycore_pyerrors.h"     // _Py_UTF8_Edit_Cost()
 #include "pycore_pystate.h"      // _PyThreadState_GET()
-#include "osdefs.h"               // MAXPATHLEN
+#include "osdefs.h"              // MAXPATHLEN
 
 
 static PyObject *
@@ -491,6 +492,38 @@ decode_locale_ex(PyObject *self, PyObject *args)
     return res;
 }
 
+static PyObject *record_list = NULL;
+
+static PyObject *
+set_eval_frame_default(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState_Get(), _PyEval_EvalFrameDefault);
+    Py_CLEAR(record_list);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+record_eval(PyThreadState *tstate, struct _PyInterpreterFrame *f, int exc)
+{
+    PyList_Append(record_list, f->f_func->func_name);
+    return _PyEval_EvalFrameDefault(tstate, f, exc);
+}
+
+
+static PyObject *
+set_eval_frame_record(PyObject *self, PyObject *list)
+{
+    if (!PyList_Check(list)) {
+        PyErr_SetString(PyExc_TypeError, "argument must be a list");
+        return NULL;
+    }
+    Py_CLEAR(record_list);
+    Py_INCREF(list);
+    record_list = list;
+    _PyInterpreterState_SetEvalFrameFunc(PyInterpreterState_Get(), record_eval);
+    Py_RETURN_NONE;
+}
+
 
 static PyMethodDef TestMethods[] = {
     {"get_configs", get_configs, METH_NOARGS},
@@ -508,6 +541,8 @@ static PyMethodDef TestMethods[] = {
     {"get_getpath_codeobject", get_getpath_codeobject, METH_NOARGS, NULL},
     {"EncodeLocaleEx", encode_locale_ex, METH_VARARGS},
     {"DecodeLocaleEx", decode_locale_ex, METH_VARARGS},
+    {"set_eval_frame_default", set_eval_frame_default, METH_NOARGS, NULL},
+    {"set_eval_frame_record", set_eval_frame_record, METH_O, NULL},
     {NULL, NULL} /* sentinel */
 };
 

@@ -4890,6 +4890,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 
         TARGET(CALL_PY_EXACT_ARGS) {
             assert(call_shape.kwnames == NULL);
+            DEOPT_IF(tstate->interp->eval_frame, CALL);
             _PyCallCache *cache = (_PyCallCache *)next_instr;
             int is_meth = is_method(stack_pointer, oparg);
             int argcount = oparg + is_meth;
@@ -4923,6 +4924,7 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 
         TARGET(CALL_PY_WITH_DEFAULTS) {
             assert(call_shape.kwnames == NULL);
+            DEOPT_IF(tstate->interp->eval_frame, CALL);
             _PyCallCache *cache = (_PyCallCache *)next_instr;
             int is_meth = is_method(stack_pointer, oparg);
             int argcount = oparg + is_meth;

diff --git a/Python/specialize.c b/Python/specialize.c
@@ -435,6 +435,7 @@ initial_counter_value(void) {
 #define SPEC_FAIL_CALL_METHOD_WRAPPER 26
 #define SPEC_FAIL_CALL_OPERATOR_WRAPPER 27
 #define SPEC_FAIL_CALL_PYFUNCTION 28
+#define SPEC_FAIL_CALL_PEP_523 29
 
 /* COMPARE_OP */
 #define SPEC_FAIL_COMPARE_OP_DIFFERENT_TYPES 12
@@ -1466,6 +1467,11 @@ specialize_py_call(PyFunctionObject *func, _Py_CODEUNIT *instr, int nargs,
     assert(_Py_OPCODE(*instr) == CALL_ADAPTIVE);
     PyCodeObject *code = (PyCodeObject *)func->func_code;
     int kind = function_kind(code);
+    /* Don't specialize if PEP 523 is active */
+    if (_PyInterpreterState_GET()->eval_frame) {
+        SPECIALIZATION_FAIL(CALL, SPEC_FAIL_CALL_PEP_523);
+        return -1;
+    }
     if (kwnames) {
         SPECIALIZATION_FAIL(CALL, SPEC_FAIL_CALL_KWNAMES);
         return -1;