gh-105927: PyImport_AddModule() uses _PyWeakref_GET_REF() (#106001)

vstinner · web-flow · commit ee52158f2098 · 2023-06-23T01:23:08.000+02:00
It now raises an exception if sys.modules doesn't hold a strong
reference to the module.

Elaborate the comment explaining why a weak reference is used to
create a borrowed reference.
diff --git 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 */
 }
 
