gh-105922: Add PyImport_AddModuleRef() function (#105923)

* Add tests on PyImport_AddModuleRef(), PyImport_AddModule() and
  PyImport_AddModuleObject().
* pythonrun.c: Replace Py_XNewRef(PyImport_AddModule(name)) with
  PyImport_AddModuleRef(name).

2 files changed, 31 insertions, 12 deletions

diff --git a/Python/import.c b/Python/import.c
index 779af55..969902a 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -351,19 +351,37 @@ import_add_module(PyThreadState *tstate, PyObject *name)
 }
 
 PyObject *
+PyImport_AddModuleRef(const char *name)
+{
+    PyObject *name_obj = PyUnicode_FromString(name);
+    if (name_obj == NULL) {
+        return NULL;
+    }
+    PyThreadState *tstate = _PyThreadState_GET();
+    PyObject *module = import_add_module(tstate, name_obj);
+    Py_DECREF(name_obj);
+    return module;
+}
+
+
+PyObject *
 PyImport_AddModuleObject(PyObject *name)
 {
     PyThreadState *tstate = _PyThreadState_GET();
     PyObject *mod = import_add_module(tstate, name);
-    if (mod) {
-        PyObject *ref = PyWeakref_NewRef(mod, NULL);
-        Py_DECREF(mod);
-        if (ref == NULL) {
-            return NULL;
-        }
-        mod = PyWeakref_GetObject(ref);
-        Py_DECREF(ref);
+    if (!mod) {
+        return NULL;
     }
+
+    // gh-86160: PyImport_AddModuleObject() returns a borrowed reference
+    PyObject *ref = PyWeakref_NewRef(mod, NULL);
+    Py_DECREF(mod);
+    if (ref == NULL) {
+        return NULL;
+    }
+
+    mod = PyWeakref_GetObject(ref);
+    Py_DECREF(ref);
     return mod; /* borrowed reference */
 }
 
@@ -2240,11 +2258,12 @@ init_importlib(PyThreadState *tstate, PyObject *sysmod)
     if (PyImport_ImportFrozenModule("_frozen_importlib") <= 0) {
         return -1;
     }
-    PyObject *importlib = PyImport_AddModule("_frozen_importlib"); // borrowed
+
+    PyObject *importlib = PyImport_AddModuleRef("_frozen_importlib");
     if (importlib == NULL) {
         return -1;
     }
-    IMPORTLIB(interp) = Py_NewRef(importlib);
+    IMPORTLIB(interp) = importlib;
 
     // Import the _imp module
     if (verbose) {
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index f9e798e..af82fa4 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -406,7 +406,7 @@ _PyRun_SimpleFileObject(FILE *fp, PyObject *filename, int closeit,
 {
     int ret = -1;
 
-    PyObject *main_module = Py_XNewRef(PyImport_AddModule("__main__"));
+    PyObject *main_module = PyImport_AddModuleRef("__main__");
     if (main_module == NULL)
         return -1;
     PyObject *dict = PyModule_GetDict(main_module);  // borrowed ref
@@ -502,7 +502,7 @@ PyRun_SimpleFileExFlags(FILE *fp, const char *filename, int closeit,
 int
 PyRun_SimpleStringFlags(const char *command, PyCompilerFlags *flags)
 {
-    PyObject *main_module = Py_XNewRef(PyImport_AddModule("__main__"));
+    PyObject *main_module = PyImport_AddModuleRef("__main__");
     if (main_module == NULL) {
         return -1;
     }