diff --git a/Python/import.c b/Python/import.c index 05da9506cc0b64..97da8bba8adff9 100644 --- a/Python/import.c +++ b/Python/import.c @@ -12,6 +12,7 @@ #include "pycore_pymem.h" // _PyMem_SetDefaultAllocator() #include "pycore_pystate.h" // _PyInterpreterState_GET() #include "pycore_sysmodule.h" // _PySys_Audit() +#include "pycore_weakref.h" // _PyWeakref_GET_REF() #include "marshal.h" // PyMarshal_ReadObjectFromString() #include "importdl.h" // _PyImport_DynLoadFiletab #include "pydtrace.h" // PyDTrace_IMPORT_FIND_LOAD_START_ENABLED() @@ -373,15 +374,30 @@ PyImport_AddModuleObject(PyObject *name) return NULL; } - // gh-86160: PyImport_AddModuleObject() returns a borrowed reference + // gh-86160: PyImport_AddModuleObject() returns a borrowed reference. + // Create a weak reference to produce a borrowed reference, since it can + // become NULL. sys.modules type can be different than dict and it is not + // guaranteed that it keeps a strong reference to the module. It can be a + // custom mapping with __getitem__() which returns a new object or removes + // returned object, or __setitem__ which does nothing. There is so much + // unknown. With weakref we can be sure that we get either a reference to + // live object or NULL. + // + // Use PyImport_AddModuleRef() to avoid these issues. PyObject *ref = PyWeakref_NewRef(mod, NULL); Py_DECREF(mod); if (ref == NULL) { return NULL; } - - mod = PyWeakref_GetObject(ref); + mod = _PyWeakref_GET_REF(ref); Py_DECREF(ref); + Py_XDECREF(mod); + + if (mod == NULL && !PyErr_Occurred()) { + PyErr_SetString(PyExc_RuntimeError, + "sys.modules does not hold a strong reference " + "to the module"); + } return mod; /* borrowed reference */ }