Use _PyImport_RunModInitFunc() in create_builtin().

Author: ericsnowcurrently

Include/internal/pycore_importdl.h

@@ -36,6 +36,9 @@ extern int _Py_ext_module_loader_info_init(
     struct _Py_ext_module_loader_info *info,
     PyObject *name,
     PyObject *filename);
+extern int _Py_ext_module_loader_info_init_for_builtin(
+    struct _Py_ext_module_loader_info *p_info,
+    PyObject *name);
 extern int _Py_ext_module_loader_info_init_from_spec(
     struct _Py_ext_module_loader_info *info,
     PyObject *spec);

Python/import.c

@@ -1472,7 +1472,7 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
     PyModuleDef *def = NULL;
 
     struct _Py_ext_module_loader_info info;
-    if (_Py_ext_module_loader_info_init(&info, name, NULL) < 0) {
+    if (_Py_ext_module_loader_info_init_for_builtin(&info, name) < 0) {
         return NULL;
     }
 
@@ -1506,26 +1506,24 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
         goto finally;
     }
 
-    mod = p0();
-    if (mod == NULL) {
+    struct _Py_ext_module_loader_result res;
+    if (_PyImport_RunModInitFunc(p0, &info, &res) < 0) {
         goto finally;
     }
 
-    if (PyObject_TypeCheck(mod, &PyModuleDef_Type)) {
-        def = (PyModuleDef*)mod;
+    mod = res.module;
+    res.module = NULL;
+    def = res.def;
+    assert(def != NULL);
+
+    if (mod == NULL) {
         assert(!is_singlephase(def));
         mod = PyModule_FromDefAndSpec(def, spec);
         if (mod == NULL) {
             goto finally;
         }
     }
     else {
-        assert(PyModule_Check(mod));
-        def = PyModule_GetDef(mod);
-        if (def == NULL) {
-            Py_CLEAR(mod);
-            goto finally;
-        }
         assert(is_singlephase(def));
 
         /* Remember pointer to module init function. */
@@ -1540,12 +1538,14 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
             singlephase.m_dict = PyModule_GetDict(mod);
             assert(singlephase.m_dict != NULL);
         }
+
         if (update_global_state_for_extension(
                 tstate, info.name, info.path, def, &singlephase) < 0)
         {
             Py_CLEAR(mod);
             goto finally;
         }
+
         PyObject *modules = get_modules_dict(tstate, true);
         if (finish_singlephase_extension(
                 tstate, mod, def, info.name, modules) < 0)
@@ -3942,7 +3942,7 @@ _imp_create_dynamic_impl(PyObject *module, PyObject *spec, PyObject *file)
         }
 
         /* Remember pointer to module init function. */
-        res.def->m_base.m_init = p0;
+        def->m_base.m_init = p0;
 
         /* Remember the filename as the __file__ attribute */
         if (PyModule_AddObjectRef(mod, "__file__", info.filename) < 0) {

Python/importdl.c

@@ -117,6 +117,7 @@ _Py_ext_module_loader_info_init(struct _Py_ext_module_loader_info *p_info,
         _Py_ext_module_loader_info_clear(&info);
         return -1;
     }
+    assert(PyUnicode_GetLength(name) > 0);
     info.name = Py_NewRef(name);
 
     info.name_encoded = get_encoded_name(info.name, &info.hook_prefix);
@@ -158,6 +159,31 @@ _Py_ext_module_loader_info_init(struct _Py_ext_module_loader_info *p_info,
     return 0;
 }
 
+int
+_Py_ext_module_loader_info_init_for_builtin(
+                            struct _Py_ext_module_loader_info *info,
+                            PyObject *name)
+{
+    assert(PyUnicode_Check(name));
+    assert(PyUnicode_FindChar(name, '.', 0, PyUnicode_GetLength(name), -1) == -1);
+    assert(PyUnicode_GetLength(name) > 0);
+
+    PyObject *name_encoded = PyUnicode_AsEncodedString(name, "ascii", NULL);
+    if (name_encoded == NULL) {
+        return -1;
+    }
+
+    *info = (struct _Py_ext_module_loader_info){
+        .name=Py_NewRef(name),
+        .name_encoded=name_encoded,
+        /* We won't need filename. */
+        .path=name,
+        .hook_prefix=ascii_only_prefix,
+        .newcontext=NULL,
+    };
+    return 0;
+}
+
 int
 _Py_ext_module_loader_info_init_from_spec(
                             struct _Py_ext_module_loader_info *p_info,
@@ -280,14 +306,20 @@ _PyImport_RunModInitFunc(PyModInitFunction p0,
         /* single-phase init (legacy) */
         res.module = m;
 
-        res.def = PyModule_GetDef(m);
-        if (res.def == NULL) {
-            PyErr_Clear();
+        if (!PyModule_Check(m)) {
             PyErr_Format(PyExc_SystemError,
                          "initialization of %s did not return an extension "
                          "module", name_buf);
             goto error;
         }
+
+        res.def = _PyModule_GetDef(m);
+        if (res.def == NULL) {
+            PyErr_Format(PyExc_SystemError,
+                         "initialization of %s did not return a valid extension "
+                         "module", name_buf);
+            goto error;
+        }
     }
 
     assert(!PyErr_Occurred());