python · vstinner · Nov 22, 2019 · Nov 22, 2019
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -147,7 +147,6 @@ struct _ts {
   The caller must hold the GIL.*/
 PyAPI_FUNC(PyInterpreterState *) _PyInterpreterState_Get(void);

-PyAPI_FUNC(int) _PyState_AddModule(PyObject*, struct PyModuleDef*);
 PyAPI_FUNC(PyThreadState *) _PyThreadState_Prealloc(PyInterpreterState *);

 /* Similar to PyThreadState_Get(), but don't issue a fatal error

diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
@@ -324,6 +324,12 @@ extern void _PyInterpreterState_ClearModules(PyInterpreterState *interp);

 PyAPI_FUNC(void) _PyGILState_Reinit(_PyRuntimeState *runtime);

+
+PyAPI_FUNC(int) _PyState_AddModule(
+    PyThreadState *tstate,
+    PyObject* module,
+    struct PyModuleDef* def);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/Python/import.c b/Python/import.c
@@ -695,50 +695,62 @@ int
 _PyImport_FixupExtensionObject(PyObject *mod, PyObject *name,
                               PyObject *filename, PyObject *modules)
 {
-    PyObject *dict, *key;
-    struct PyModuleDef *def;
-    int res;
-    if (extensions == NULL) {
-        extensions = PyDict_New();
-        if (extensions == NULL)
-            return -1;
-    }
    if (mod == NULL || !PyModule_Check(mod)) {
        PyErr_BadInternalCall();
        return -1;
    }
-    def = PyModule_GetDef(mod);
+
+    struct PyModuleDef *def = PyModule_GetDef(mod);
    if (!def) {
        PyErr_BadInternalCall();
        return -1;
    }
-    if (PyObject_SetItem(modules, name, mod) < 0)
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    if (PyObject_SetItem(modules, name, mod) < 0) {
        return -1;
-    if (_PyState_AddModule(mod, def) < 0) {
+    }
+    if (_PyState_AddModule(tstate, mod, def) < 0) {
        PyMapping_DelItem(modules, name);
        return -1;
    }
-    if (def->m_size == -1) {
-        if (def->m_base.m_copy) {
-            /* Somebody already imported the module,
-               likely under a different name.
-               XXX this should really not happen. */
-            Py_CLEAR(def->m_base.m_copy);
+
+    if (_Py_IsMainInterpreter(tstate)) {
+        if (def->m_size == -1) {
+            if (def->m_base.m_copy) {
+                /* Somebody already imported the module,
+                   likely under a different name.
+                   XXX this should really not happen. */
+                Py_CLEAR(def->m_base.m_copy);
+            }
+            PyObject *dict = PyModule_GetDict(mod);
+            if (dict == NULL) {
+                return -1;
+            }
+            def->m_base.m_copy = PyDict_Copy(dict);
+            if (def->m_base.m_copy == NULL) {
+                return -1;
+            }
        }
-        dict = PyModule_GetDict(mod);
-        if (dict == NULL)
+
+        if (extensions == NULL) {
+            extensions = PyDict_New();
+            if (extensions == NULL) {
+                return -1;
+            }
+        }
+
+        PyObject *key = PyTuple_Pack(2, filename, name);
+        if (key == NULL) {
            return -1;
-        def->m_base.m_copy = PyDict_Copy(dict);
-        if (def->m_base.m_copy == NULL)
+        }
+        int res = PyDict_SetItem(extensions, key, (PyObject *)def);
+        Py_DECREF(key);
+        if (res < 0) {
            return -1;
+        }
    }
-    key = PyTuple_Pack(2, filename, name);
-    if (key == NULL)
-        return -1;
-    res = PyDict_SetItem(extensions, key, (PyObject *)def);
-    Py_DECREF(key);
-    if (res < 0)
-        return -1;
+
    return 0;
 }

@@ -801,7 +813,7 @@ import_find_extension(PyThreadState *tstate, PyObject *name,
        }
        Py_DECREF(mod);
    }
-    if (_PyState_AddModule(mod, def) < 0) {
+    if (_PyState_AddModule(tstate, mod, def) < 0) {
        PyMapping_DelItem(modules, name);
        return NULL;
    }

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -661,9 +661,8 @@ PyState_FindModule(struct PyModuleDef* module)
 }

 int
-_PyState_AddModule(PyObject* module, struct PyModuleDef* def)
+_PyState_AddModule(PyThreadState *tstate, PyObject* module, struct PyModuleDef* def)
 {
-    PyInterpreterState *state;
    if (!def) {
        assert(PyErr_Occurred());
        return -1;
@@ -673,37 +672,45 @@ _PyState_AddModule(PyObject* module, struct PyModuleDef* def)
                        "PyState_AddModule called on module with slots");
        return -1;
    }
-    state = _PyInterpreterState_GET_UNSAFE();
-    if (!state->modules_by_index) {
-        state->modules_by_index = PyList_New(0);
-        if (!state->modules_by_index)
+
+    PyInterpreterState *interp = tstate->interp;
+    if (!interp->modules_by_index) {
+        interp->modules_by_index = PyList_New(0);
+        if (!interp->modules_by_index) {
            return -1;
+        }
    }
-    while (PyList_GET_SIZE(state->modules_by_index) <= def->m_base.m_index)
-        if (PyList_Append(state->modules_by_index, Py_None) < 0)
+
+    while (PyList_GET_SIZE(interp->modules_by_index) <= def->m_base.m_index) {
+        if (PyList_Append(interp->modules_by_index, Py_None) < 0) {
            return -1;
+        }
+    }
+
    Py_INCREF(module);
-    return PyList_SetItem(state->modules_by_index,
+    return PyList_SetItem(interp->modules_by_index,
                          def->m_base.m_index, module);
 }

 int
 PyState_AddModule(PyObject* module, struct PyModuleDef* def)
 {
-    Py_ssize_t index;
-    PyInterpreterState *state = _PyInterpreterState_GET_UNSAFE();
    if (!def) {
        Py_FatalError("PyState_AddModule: Module Definition is NULL");
        return -1;
    }
-    index = def->m_base.m_index;
-    if (state->modules_by_index &&
-        index < PyList_GET_SIZE(state->modules_by_index) &&
-        module == PyList_GET_ITEM(state->modules_by_index, index)) {
+
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyInterpreterState *interp = tstate->interp;
+    Py_ssize_t index = def->m_base.m_index;
+    if (interp->modules_by_index &&
+        index < PyList_GET_SIZE(interp->modules_by_index) &&
+        module == PyList_GET_ITEM(interp->modules_by_index, index))
+    {
        Py_FatalError("PyState_AddModule: Module already added!");
        return -1;
    }
-    return _PyState_AddModule(module, def);
+    return _PyState_AddModule(tstate, module, def);
 }

 int