gh-91053: Add an optional callback that is invoked whenever a function is modified

Include/cpython/funcobject.h

@@ -131,6 +131,44 @@ PyAPI_DATA(PyTypeObject) PyStaticMethod_Type;
 PyAPI_FUNC(PyObject *) PyClassMethod_New(PyObject *);
 PyAPI_FUNC(PyObject *) PyStaticMethod_New(PyObject *);
 
+#define FOREACH_FUNC_EVENT(V) \
+ V(CREATED) \
+ V(DESTROY) \
+ V(MODIFY_CODE) \
+ V(MODIFY_DEFAULTS) \
+ V(MODIFY_KWDEFAULTS)
+
+typedef enum {
+ #define DEF_EVENT(EVENT) PYFUNC_EVENT_##EVENT,
+ FOREACH_FUNC_EVENT(DEF_EVENT)
+ #undef DEF_EVENT
+} PyFunction_Event;
+
+/*
+ * A callback that is invoked for different events in a function's lifecycle.
+ *
+ * The callback is invoked with a borrowed reference to func, after it is
+ * created and before it is modified or destroyed. The callback should not
+ * modify func.
+ *
+ * When a function's code object, defaults, or kwdefaults are modified the
+ * callback will be invoked with the respective event and new_value will
+ * contain a borrowed reference to the new value that is about to be stored in
+ * the function. Otherwise the third argument is NULL.
+ */
+typedef void(*PyFunction_EventCallback)(
+ PyFunction_Event event,
+ PyFunctionObject *func,
+ PyObject *new_value);
+
+/*
+ * Set the callback that will be invoked for function lifecycle events.
+ *
+ * Pass NULL to clear the callback.
+ */
+PyAPI_FUNC(void) PyFunction_SetEventCallback(PyFunction_EventCallback callback);
+PyAPI_FUNC(PyFunction_EventCallback) PyFunction_GetEventCallback(void);
+
 #ifdef __cplusplus
 }
 #endif

Lib/test/test_func_events.py

@@ -0,0 +1,44 @@
+import unittest
+from _testcapi import (
+ PYFUNC_EVENT_CREATED,
+ PYFUNC_EVENT_DESTROY,
+ PYFUNC_EVENT_MODIFY_CODE,
+ PYFUNC_EVENT_MODIFY_DEFAULTS,
+ PYFUNC_EVENT_MODIFY_KWDEFAULTS,
+ restore_func_event_callback,
+ set_func_event_callback,
+)
+
+
+class FuncEventsTest(unittest.TestCase):
+ def test_func_events_dispatched(self):
+ event = None
+ def handle_func_event(*args):
+ nonlocal event
+ event = args
+ set_func_event_callback(handle_func_event)
+
+ try:
+ def myfunc():
+ pass
+ self.assertEqual(event, (PYFUNC_EVENT_CREATED, myfunc, None))
+ myfunc_id = id(myfunc)
+
+ new_code = self.test_func_events_dispatched.__code__
+ myfunc.__code__ = new_code
+ self.assertEqual(event, (PYFUNC_EVENT_MODIFY_CODE, myfunc, new_code))
+
+ new_defaults = (123,)
+ myfunc.__defaults__ = new_defaults
+ self.assertEqual(event, (PYFUNC_EVENT_MODIFY_DEFAULTS, myfunc, new_defaults))
+
+ new_kwdefaults = {"self": 123}
+ myfunc.__kwdefaults__ = new_kwdefaults
+ self.assertEqual(event, (PYFUNC_EVENT_MODIFY_KWDEFAULTS, myfunc, new_kwdefaults))
+
+ # Clear event's reference to func
+ event = None
+ del myfunc
+ self.assertEqual(event, (PYFUNC_EVENT_DESTROY, myfunc_id, None))
+ finally:
+ restore_func_event_callback()

Modules/Setup.stdlib.in

@@ -169,7 +169,7 @@
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
 @MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c
-@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c
+@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/unicode.c _testcapi/func_events.c
 
 # Some testing modules MUST be built as shared libraries.
 *shared*

Modules/_testcapi/func_events.c

@@ -0,0 +1,100 @@
+#include "parts.h"
+
+static PyObject *pyfunc_callback = NULL;
+static PyFunction_EventCallback orig_callback = NULL;
+
+static void
+call_pyfunc_callback(PyFunction_Event event, PyFunctionObject *func, PyObject *new_value)
+{
+ PyObject *event_obj = PyLong_FromLong(event);
+ if (event_obj == NULL) {
+ PyErr_Clear();
+ return;
+ }
+ if (new_value == NULL) {
+ new_value = Py_None;
+ }
+ Py_INCREF(new_value);
+ /* Don't expose a function that's about to be destroyed to managed code */
+ PyObject *func_or_id = (PyObject *) func;
+ if (event == PYFUNC_EVENT_DESTROY) {
+ func_or_id = PyLong_FromLong((long) func);
+ if (func_or_id == NULL) {
+ Py_DECREF(new_value);
+ Py_DECREF(event_obj);
+ return;
+ }
+ } else {
+ Py_INCREF(func);
+ }
+ PyObject *stack[] = {event_obj, func_or_id, new_value};
+ PyObject_Vectorcall(pyfunc_callback, stack, 3, NULL);
+ Py_DECREF(new_value);
+ Py_DECREF(event_obj);
+ Py_DECREF(func_or_id);
+}
+
+static int
+add_event(PyObject *module, const char *name, PyFunction_Event event)
+{
+ PyObject *value = PyLong_FromLong(event);
+ if (value == NULL) {
+ return -1;
+ }
+ int ok = PyModule_AddObjectRef(module, name, value);
+ Py_DECREF(value);
+ return ok;
+}
+
+static PyObject *
+set_func_event_callback(PyObject *self, PyObject *func)
+{
+ if (!PyFunction_Check(func)) {
+ PyErr_SetString(PyExc_TypeError, "'func' must be a function");
+ return NULL;
+ }
+ if (pyfunc_callback != NULL) {
+ PyErr_SetString(PyExc_RuntimeError, "already set callback");
+ return NULL;
+ }
+ Py_INCREF(func);
+ pyfunc_callback = func;
+ orig_callback = PyFunction_GetEventCallback();
+ PyFunction_SetEventCallback(call_pyfunc_callback);
+ Py_RETURN_NONE;
+}
+
+static PyObject *
+restore_func_event_callback() {
+ if (pyfunc_callback == NULL) {
+ PyErr_SetString(PyExc_RuntimeError, "nothing to restore");
+ return NULL;
+ }
+ PyFunction_SetEventCallback(orig_callback);
+ orig_callback = NULL;
+ Py_CLEAR(pyfunc_callback);
+ Py_RETURN_NONE;
+}
+
+static PyMethodDef TestMethods[] = {
+ {"set_func_event_callback", set_func_event_callback, METH_O},
+ {"restore_func_event_callback", restore_func_event_callback, METH_NOARGS},
+ {NULL},
+};
+
+int
+_PyTestCapi_Init_FuncEvents(PyObject *m) {
+ if (PyModule_AddFunctions(m, TestMethods) < 0) {
+ return -1;
+ }
+
+ /* Expose each event as an attribute on the module */
+#define ADD_EVENT(event) \
+ if (add_event(m, "PYFUNC_EVENT_" #event, PYFUNC_EVENT_##event)) { \
+ return -1; \
+ }
+ FOREACH_FUNC_EVENT(ADD_EVENT);
+#undef ADD_EVENT
+
+ return 0;
+}

Modules/_testcapi/parts.h

@@ -27,6 +27,7 @@
 int _PyTestCapi_Init_Vectorcall(PyObject *module);
 int _PyTestCapi_Init_Heaptype(PyObject *module);
 int _PyTestCapi_Init_Unicode(PyObject *module);
+int _PyTestCapi_Init_FuncEvents(PyObject *module);
 
 #ifdef LIMITED_API_AVAILABLE
 int _PyTestCapi_Init_VectorcallLimited(PyObject *module);

Modules/_testcapimodule.c

@@ -6578,6 +6578,9 @@ PyInit__testcapi(void)
  if (_PyTestCapi_Init_Unicode(m) < 0) {
  return NULL;
  }
+ if (_PyTestCapi_Init_FuncEvents(m) < 0) {
+ return NULL;
+ }
 
 #ifndef LIMITED_API_AVAILABLE
  PyModule_AddObjectRef(m, "LIMITED_API_AVAILABLE", Py_False);

Objects/funcobject.c

@@ -9,6 +9,29 @@
 
 static uint32_t next_func_version = 1;
 
+static PyFunction_EventCallback func_event_callback = NULL;
+
+static void
+handle_func_event(PyFunction_Event event, PyFunctionObject *func, PyObject *new_value)
+{
+ if (func_event_callback == NULL) {
+ return;
+ }
+ func_event_callback(event, func, new_value);
+}
+
+void
+PyFunction_SetEventCallback(PyFunction_EventCallback callback)
+{
+ func_event_callback = callback;
+}
+
+PyFunction_EventCallback
+PyFunction_GetEventCallback()
+{
+ return func_event_callback;
+}
+
 PyFunctionObject *
 _PyFunction_FromConstructor(PyFrameConstructor *constr)
 {
@@ -40,6 +63,7 @@ _PyFunction_FromConstructor(PyFrameConstructor *constr)
  op->vectorcall = _PyFunction_Vectorcall;
  op->func_version = 0;
  _PyObject_GC_TRACK(op);
+ handle_func_event(PYFUNC_EVENT_CREATED, op, NULL);
  return op;
 }
 
@@ -116,6 +140,7 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
  op->vectorcall = _PyFunction_Vectorcall;
  op->func_version = 0;
  _PyObject_GC_TRACK(op);
+ handle_func_event(PYFUNC_EVENT_CREATED, op, NULL);
  return (PyObject *)op;
 
 error:
@@ -402,6 +427,7 @@ func_set_code(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignored))
  nclosure, nfree);
  return -1;
  }
+ handle_func_event(PYFUNC_EVENT_MODIFY_CODE, op, value);
  op->func_version = 0;
  Py_INCREF(value);
  Py_XSETREF(op->func_code, value);
@@ -487,6 +513,7 @@ func_set_defaults(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignored
  return -1;
  }
 
+ handle_func_event(PYFUNC_EVENT_MODIFY_DEFAULTS, op, value);
  op->func_version = 0;
  Py_XINCREF(value);
  Py_XSETREF(op->func_defaults, value);
@@ -529,6 +556,7 @@ func_set_kwdefaults(PyFunctionObject *op, PyObject *value, void *Py_UNUSED(ignor
  return -1;
  }
 
+ handle_func_event(PYFUNC_EVENT_MODIFY_KWDEFAULTS, op, value);
  op->func_version = 0;
  Py_XINCREF(value);
  Py_XSETREF(op->func_kwdefaults, value);
@@ -712,6 +740,7 @@ func_clear(PyFunctionObject *op)
 static void
 func_dealloc(PyFunctionObject *op)
 {
+ handle_func_event(PYFUNC_EVENT_DESTROY, op, NULL);
  _PyObject_GC_UNTRACK(op);
  if (op->func_weakreflist != NULL) {
  PyObject_ClearWeakRefs((PyObject *) op);