gh-148014: Accept a function name in -X presite option (#148015)

Python/pylifecycle.c

@@ -1218,6 +1218,54 @@ pyinit_main_reconfigure(PyThreadState *tstate)
 
 
 #ifdef Py_DEBUG
+// Equivalent to the Python code:
+//
+//     for part in attr.split('.'):
+//         obj = getattr(obj, part)
+static PyObject*
+presite_resolve_name(PyObject *obj, PyObject *attr)
+{
+    obj = Py_NewRef(obj);
+    attr = Py_NewRef(attr);
+    PyObject *res;
+
+    while (1) {
+        Py_ssize_t len = PyUnicode_GET_LENGTH(attr);
+        Py_ssize_t pos = PyUnicode_FindChar(attr, '.', 0, len, 1);
+        if (pos < 0) {
+            break;
+        }
+
+        PyObject *name = PyUnicode_Substring(attr, 0, pos);
+        if (name == NULL) {
+            goto error;
+        }
+        res = PyObject_GetAttr(obj, name);
+        Py_DECREF(name);
+        if (res == NULL) {
+            goto error;
+        }
+        Py_SETREF(obj, res);
+
+        PyObject *suffix = PyUnicode_Substring(attr, pos + 1, len);
+        if (suffix == NULL) {
+            goto error;
+        }
+        Py_SETREF(attr, suffix);
+    }
+
+    res = PyObject_GetAttr(obj, attr);
+    Py_DECREF(obj);
+    Py_DECREF(attr);
+    return res;
+
+error:
+    Py_DECREF(obj);
+    Py_DECREF(attr);
+    return NULL;
+}
+
+
 static void
 run_presite(PyThreadState *tstate)
 {
@@ -1228,22 +1276,68 @@ run_presite(PyThreadState *tstate)
         return;
     }
 
-    PyObject *presite_modname = PyUnicode_FromWideChar(
-        config->run_presite,
-        wcslen(config->run_presite)
-    );
-    if (presite_modname == NULL) {
-        fprintf(stderr, "Could not convert pre-site module name to unicode\n");
+    PyObject *presite = PyUnicode_FromWideChar(config->run_presite, -1);
+    if (presite == NULL) {
+        fprintf(stderr, "Could not convert pre-site command to Unicode\n");
+        _PyErr_Print(tstate);
+        return;
+    }
+
+    // Accept "mod_name" and "mod_name:func_name" entry point syntax
+    Py_ssize_t len = PyUnicode_GET_LENGTH(presite);
+    Py_ssize_t pos = PyUnicode_FindChar(presite, ':', 0, len, 1);
+    PyObject *mod_name = NULL;
+    PyObject *func_name = NULL;
+    PyObject *module = NULL;
+    if (pos > 0) {
+        mod_name = PyUnicode_Substring(presite, 0, pos);
+        if (mod_name == NULL) {
+            goto error;
+        }
+
+        func_name = PyUnicode_Substring(presite, pos + 1, len);
+        if (func_name == NULL) {
+            goto error;
+        }
     }
     else {
-        PyObject *presite = PyImport_Import(presite_modname);
-        if (presite == NULL) {
-            fprintf(stderr, "pre-site import failed:\n");
-            _PyErr_Print(tstate);
-        }
-        Py_XDECREF(presite);
-        Py_DECREF(presite_modname);
+        mod_name = Py_NewRef(presite);
     }
+
+    // mod_name can contain dots (ex: "math.integer")
+    module = PyImport_Import(mod_name);
+    if (module == NULL) {
+        goto error;
+    }
+
+    if (func_name != NULL) {
+        PyObject *func = presite_resolve_name(module, func_name);
+        if (func == NULL) {
+            goto error;
+        }
+
+        PyObject *res = PyObject_CallNoArgs(func);
+        Py_DECREF(func);
+        if (res == NULL) {
+            goto error;
+        }
+        Py_DECREF(res);
+    }
+
+    Py_DECREF(presite);
+    Py_DECREF(mod_name);
+    Py_XDECREF(func_name);
+    Py_DECREF(module);
+    return;
+
+error:
+    fprintf(stderr, "pre-site failed:\n");
+    _PyErr_Print(tstate);
+
+    Py_DECREF(presite);
+    Py_XDECREF(mod_name);
+    Py_XDECREF(func_name);
+    Py_XDECREF(module);
 }
 #endif