pythongh-142417: Add PyInitConfig_SetInitCallback() function

Add init_callback and init_callback_arg members to PyConfig.

Author: vstinner

Doc/c-api/init_config.rst

@@ -238,6 +238,22 @@ Module
    Similar to the :c:func:`PyImport_AppendInittab` function.
 
 
+Initialization Callback
+-----------------------
+
+.. c:function:: int PyInitConfig_SetInitCallback(PyInitConfig *config, PyStatus (*callback)(void *arg), void *arg)
+
+   Set an initialization callback. It is called once the Python interpreter is
+   initialized and before the first import.
+
+   A single callback can be registered. If this function is called more than
+   once, the previous callback is overridden.
+
+   * Return ``0`` on success.
+   * Set an error in *config* and return ``-1`` on error.
+
+   .. versionadded:: next
+
 Initialize Python
 -----------------
 

Include/cpython/initconfig.h

@@ -215,6 +215,10 @@ typedef struct PyConfig {
     wchar_t *run_module;
     wchar_t *run_filename;
 
+    /* --- Initialization callback ------------------- */
+    PyStatus (*init_callback)(void *arg);
+    void *init_callback_arg;
+
     /* --- Set by Py_Main() -------------------------- */
     wchar_t *sys_path_0;
 
@@ -323,6 +327,10 @@ PyAPI_FUNC(int) PyInitConfig_AddModule(PyInitConfig *config,
     const char *name,
     PyObject* (*initfunc)(void));
 
+PyAPI_FUNC(int) PyInitConfig_SetInitCallback(PyInitConfig *config,
+    PyStatus (*callback)(void *arg),
+    void *arg);
+
 PyAPI_FUNC(int) Py_InitializeFromInitConfig(PyInitConfig *config);
 
 

Lib/test/test_embed.py

@@ -1885,6 +1885,16 @@ def test_init_in_background_thread(self):
         out, err = self.run_embedded_interpreter("test_init_in_background_thread")
         self.assertEqual(err, "")
 
+    def test_init_callback(self):
+        out, err = self.run_embedded_interpreter("test_init_callback")
+        modules = [
+            '_frozen_importlib', '_imp', '_thread', '_warnings', '_weakref',
+            'builtins', 'sys']
+        self.assertEqual(out.splitlines(),
+            ["Hello Callback!",
+             f"sys.modules: {modules}"])
+        self.assertEqual(err, "")
+
 
 class AuditingTests(EmbeddingTestsMixin, unittest.TestCase):
     def test_open_code_hook(self):

Programs/_testembed.c

@@ -1741,6 +1741,16 @@ static int initconfig_getint(PyInitConfig *config, const char *name)
     return (int)value;
 }
 
+static void initconfig_error(PyInitConfig *config)
+{
+    const char *err_msg;
+    int res = PyInitConfig_GetError(config, &err_msg);
+    assert(res == 1);
+
+    printf("Python init failed: %s\n", err_msg);
+    PyInitConfig_Free(config);
+}
+
 
 static int test_initconfig_api(void)
 {
@@ -1795,12 +1805,8 @@ static int test_initconfig_api(void)
     return 0;
 
 error:
-    {
-        const char *err_msg;
-        (void)PyInitConfig_GetError(config, &err_msg);
-        printf("Python init failed: %s\n", err_msg);
-        exit(1);
-    }
+    initconfig_error(config);
+    return 1;
 }
 
 
@@ -1953,12 +1959,8 @@ static int test_initconfig_module(void)
     return 0;
 
 error:
-    {
-        const char *err_msg;
-        (void)PyInitConfig_GetError(config, &err_msg);
-        printf("Python init failed: %s\n", err_msg);
-        exit(1);
-    }
+    initconfig_error(config);
+    return 1;
 }
 
 
@@ -2170,6 +2172,77 @@ static int test_init_in_background_thread(void)
 }
 
 
+static PyStatus init_callback(void *arg)
+{
+    const char *msg = (const char*)arg;
+    printf("%s\n", msg);
+
+    PyObject *modules = PySys_GetAttrString("modules");
+    if (modules == NULL) {
+        return PyStatus_Error("failed to get sys.modules");
+    }
+
+    PyObject *builtins = PyEval_GetBuiltins();  // borrowed ref
+    if (builtins == NULL) {
+        Py_DECREF(modules);
+        return PyStatus_Error("failed to get builtins");
+    }
+
+    PyObject *sorted;
+    if (PyDict_GetItemStringRef(builtins, "sorted", &sorted) <= 0) {
+        Py_DECREF(modules);
+        return PyStatus_Error("failed to get sorted");
+    }
+
+    PyObject *names = PyObject_CallOneArg(sorted, modules);
+    Py_DECREF(modules);
+    if (names == NULL) {
+        return PyStatus_Error("sorted failed");
+    }
+
+    PySys_FormatStdout("sys.modules: %R\n", names);
+    Py_DECREF(names);
+
+    return PyStatus_Ok();
+}
+
+
+static int test_init_callback(void)
+{
+    PyInitConfig *config = PyInitConfig_Create();
+    if (config == NULL) {
+        printf("Init allocation error\n");
+        return 1;
+    }
+
+    if (PyInitConfig_SetStr(config, "program_name", PROGRAM_NAME_UTF8) < 0) {
+        goto error;
+    }
+
+    const char *ignored_msg = "ignored_msg";
+    if (PyInitConfig_SetInitCallback(config, init_callback, (void*)ignored_msg) < 0) {
+        goto error;
+    }
+
+    // PyInitConfig_SetInitCallback() can be called more than once, but the
+    // previous callback and callback argument are overridden.
+    const char *msg = "Hello Callback!";
+    if (PyInitConfig_SetInitCallback(config, init_callback, (void*)msg) < 0) {
+        goto error;
+    }
+
+    if (Py_InitializeFromInitConfig(config) < 0) {
+        goto error;
+    }
+    PyInitConfig_Free(config);
+    return 0;
+
+error:
+    initconfig_error(config);
+    return 1;
+}
+
+
 #ifndef MS_WINDOWS
 #include "test_frozenmain.h"      // M_test_frozenmain
 
@@ -2658,6 +2731,7 @@ static struct TestCase TestCases[] = {
     {"test_init_use_frozen_modules", test_init_use_frozen_modules},
     {"test_init_main_interpreter_settings", test_init_main_interpreter_settings},
     {"test_init_in_background_thread", test_init_in_background_thread},
+    {"test_init_callback", test_init_callback},
 
     // Audit
     {"test_open_code_hook", test_open_code_hook},

Python/initconfig.c

@@ -1274,6 +1274,10 @@ _PyConfig_Copy(PyConfig *config, const PyConfig *config2)
             Py_UNREACHABLE();
         }
     }
+
+    config->init_callback = config2->init_callback;
+    config->init_callback_arg = config2->init_callback_arg;
+
     return _PyStatus_OK();
 }
 
@@ -4240,6 +4244,16 @@ Py_InitializeFromInitConfig(PyInitConfig *config)
 }
 
 
+int
+PyInitConfig_SetInitCallback(PyInitConfig *config,
+                             PyStatus (*callback)(void *arg), void *arg)
+{
+    config->config.init_callback = callback;
+    config->config.init_callback_arg = arg;
+    return 0;
+}
+
+
 // --- PyConfig_Get() -------------------------------------------------------
 
 static const PyConfigSpec*

Python/pylifecycle.c

@@ -1457,6 +1457,13 @@ Py_InitializeFromConfig(const PyConfig *config)
     }
     config = _PyInterpreterState_GetConfig(tstate->interp);
 
+    if (config->init_callback != NULL) {
+        status = config->init_callback(config->init_callback_arg);
+        if (_PyStatus_EXCEPTION(status)) {
+            return status;
+        }
+    }
+
     if (config->_init_main) {
         status = pyinit_main(tstate);
         if (_PyStatus_EXCEPTION(status)) {