gh-106915: Add PyImport_ImportOrAddModule() function

vstinner · vstinner · commit 2f3cc6738a34 · 2023-11-07T10:21:27.000+01:00
Remove PyImport_AddModuleRef() function.
diff --git a/Doc/c-api/import.rst b/Doc/c-api/import.rst
@@ -98,16 +98,22 @@ Importing Modules
    an exception set on failure (the module still exists in this case).
 
 
-.. c:function:: PyObject* PyImport_AddModuleRef(const char *name)
+.. c:function:: int PyImport_ImportOrAddModule(const char *name, PyObject **module)
 
-   Return the module object corresponding to a module name.
+   Create a new module and store it in :data:`sys.modules`, or get an already
+   imported module from :data:`sys.modules`.
 
-   The *name* argument may be of the form ``package.module``. First check the
-   modules dictionary if there's one there, and if not, create a new one and
-   insert it in the modules dictionary.
+   First check the modules dictionary if there's one there, and if not, create
+   a new one and insert it in the modules dictionary.
 
-   Return a :term:`strong reference` to the module on success. Return ``NULL``
-   with an exception set on failure.
+   The *name* argument may be of the form ``package.module``.
+
+   - If the module does not exist, create a module, store it in
+     :data:`sys.modules`, set *\*module* to a :term:`strong reference` to the
+     module, and return 1.
+   - If the module was already imported, set *\*module* to a :term:`strong
+     reference` to the existing module, and return 0.
+   - On error, raise an exception, set *\*module* to NULL, and return -1.
 
    The module name *name* is decoded from UTF-8.
 
@@ -122,16 +128,18 @@ Importing Modules
 
 .. c:function:: PyObject* PyImport_AddModuleObject(PyObject *name)
 
-   Similar to :c:func:`PyImport_AddModuleRef`, but return a :term:`borrowed
-   reference` and *name* is a Python :class:`str` object.
+   Similar to :c:func:`PyImport_ImportOrAddModule`, but return a :term:`borrowed
+   reference`, *name* is a Python :class:`str` object, and don't provide the
+   information if the module was created or was already imported.
 
    .. versionadded:: 3.3
 
 
 .. c:function:: PyObject* PyImport_AddModule(const char *name)
 
-   Similar to :c:func:`PyImport_AddModuleRef`, but return a :term:`borrowed
-   reference`.
+   Similar to :c:func:`PyImport_ImportOrAddModule`, but return a :term:`borrowed
+   reference`, and don't provide the information if the module was created or
+   was already imported.
 
 
 .. c:function:: PyObject* PyImport_ExecCodeModule(const char *name, PyObject *co)
diff --git a/Doc/data/refcounts.dat b/Doc/data/refcounts.dat
@@ -974,9 +974,6 @@ PyCoro_New:PyFrameObject*:frame:0:
 PyCoro_New:PyObject*:name:0:
 PyCoro_New:PyObject*:qualname:0:
 
-PyImport_AddModuleRef:PyObject*::+1:
-PyImport_AddModuleRef:const char*:name::
-
 PyImport_AddModule:PyObject*::0:reference borrowed from sys.modules
 PyImport_AddModule:const char*:name::
 
diff --git a/Doc/data/stable_abi.dat b/Doc/data/stable_abi.dat
diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst
@@ -1018,9 +1018,10 @@ New Features
   APIs accepting the format codes always use ``Py_ssize_t`` for ``#`` formats.
   (Contributed by Inada Naoki in :gh:`104922`.)
 
-* Add :c:func:`PyImport_AddModuleRef`: similar to
-  :c:func:`PyImport_AddModule`, but return a :term:`strong reference` instead
-  of a :term:`borrowed reference`.
+* Add :c:func:`PyImport_ImportOrAddModule`: similar to
+  :c:func:`PyImport_AddModule`, but get a :term:`strong reference` instead
+  of a :term:`borrowed reference` and return 0 if the module was already
+  imported.
   (Contributed by Victor Stinner in :gh:`105922`.)
 
 * Add :c:func:`PyWeakref_GetRef` function: similar to
diff --git a/Include/import.h b/Include/import.h
@@ -43,11 +43,22 @@ PyAPI_FUNC(PyObject *) PyImport_AddModuleObject(
 PyAPI_FUNC(PyObject *) PyImport_AddModule(
     const char *name            /* UTF-8 encoded string */
     );
+
 #if !defined(Py_LIMITED_API) || Py_LIMITED_API+0 >= 0x030d0000
-PyAPI_FUNC(PyObject *) PyImport_AddModuleRef(
-    const char *name            /* UTF-8 encoded string */
+// Create a new module and store it in sys.modules, or get an already imported
+// module from sys.modules.
+//
+// - If the module does not exist, create a module, store it in sys.modules,
+//   set '*module' to a strong reference to the module, and return 1.
+// - If the module was already imported, set '*module' to a strong reference
+//   to the existing module, and return 0.
+// - On error, raise an exception, set '*module' to NULL, and return -1.
+PyAPI_FUNC(int) PyImport_ImportOrAddModule(
+    const char *name,  // UTF-8 encoded string
+    PyObject **module
     );
 #endif
+
 PyAPI_FUNC(PyObject *) PyImport_ImportModule(
     const char *name            /* UTF-8 encoded string */
     );
diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
@@ -2669,24 +2669,25 @@ def test_basic_multiple_interpreters_reset_each(self):
 @cpython_only
 class CAPITests(unittest.TestCase):
     def test_pyimport_addmodule(self):
-        # gh-105922: Test PyImport_AddModuleRef(), PyImport_AddModule()
-        # and PyImport_AddModuleObject()
+        # gh-105922: Test PyImport_ImportOrAddModule(), PyImport_AddModule()
+        # and PyImport_AddModuleObject(): module already exists.
         import _testcapi
         for name in (
             'sys',     # frozen module
             'test',    # package
             __name__,  # package.module
         ):
-            _testcapi.check_pyimport_addmodule(name)
+            self.assertIn(name, sys.modules)
+            _testcapi.check_pyimport_addmodule(name, False)
 
     def test_pyimport_addmodule_create(self):
-        # gh-105922: Test PyImport_AddModuleRef(), create a new module
+        # gh-105922: Test PyImport_ImportOrAddModule(): create a new module
         import _testcapi
         name = 'dontexist'
         self.assertNotIn(name, sys.modules)
         self.addCleanup(unload, name)
 
-        mod = _testcapi.check_pyimport_addmodule(name)
+        mod = _testcapi.check_pyimport_addmodule(name, True)
         self.assertIs(mod, sys.modules[name])
 
 
diff --git a/Lib/test/test_stable_abi_ctypes.py b/Lib/test/test_stable_abi_ctypes.py
diff --git a/Misc/NEWS.d/next/C API/2023-10-31-17-40-34.gh-issue-106915.XVWI8X.rst b/Misc/NEWS.d/next/C API/2023-10-31-17-40-34.gh-issue-106915.XVWI8X.rst
@@ -0,0 +1,2 @@
+Add :c:func:`PyImport_ImportOrAddModule` function and remove
+:c:func:`PyImport_AddModuleRef` function. Patch by Victor Stinner.
diff --git a/Misc/stable_abi.toml b/Misc/stable_abi.toml
@@ -2428,7 +2428,7 @@
     added = '3.12'
 [const.Py_TPFLAGS_ITEMS_AT_END]
     added = '3.12'
-[function.PyImport_AddModuleRef]
+[function.PyImport_ImportOrAddModule]
     added = '3.13'
 [function.PyWeakref_GetRef]
     added = '3.13'
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -3036,15 +3036,19 @@ static PyObject *
 check_pyimport_addmodule(PyObject *self, PyObject *args)
 {
     const char *name;
-    if (!PyArg_ParseTuple(args, "s", &name)) {
+    int is_new;
+    if (!PyArg_ParseTuple(args, "si", &name, &is_new)) {
         return NULL;
     }
+    // name must be the name of a module which is already in sys.modules
 
-    // test PyImport_AddModuleRef()
-    PyObject *module = PyImport_AddModuleRef(name);
-    if (module == NULL) {
+    // test PyImport_ImportOrAddModule()
+    PyObject *module = UNINITIALIZED_PTR;
+    int res = PyImport_ImportOrAddModule(name, &module);
+    if (res < 0) {
         return NULL;
     }
+    assert(res == is_new);
     assert(PyModule_Check(module));
     // module is a strong reference
 
diff --git a/PC/python3dll.c b/PC/python3dll.c
diff --git a/Python/import.c b/Python/import.c
@@ -291,7 +291,7 @@ PyImport_GetModule(PyObject *name)
    if not, create a new one and insert it in the modules dictionary. */
 
 static PyObject *
-import_add_module(PyThreadState *tstate, PyObject *name)
+import_add_module(PyThreadState *tstate, PyObject *name, int *is_new)
 {
     PyObject *modules = MODULES(tstate->interp);
     if (modules == NULL) {
@@ -305,39 +305,55 @@ import_add_module(PyThreadState *tstate, PyObject *name)
         return NULL;
     }
     if (m != NULL && PyModule_Check(m)) {
+        if (is_new) {
+            *is_new = 0;
+        }
         return m;
     }
     Py_XDECREF(m);
+
     m = PyModule_NewObject(name);
-    if (m == NULL)
+    if (m == NULL) {
         return NULL;
+    }
     if (PyObject_SetItem(modules, name, m) != 0) {
         Py_DECREF(m);
         return NULL;
     }
 
+    if (is_new) {
+        *is_new = 1;
+    }
     return m;
 }
 
-PyObject *
-PyImport_AddModuleRef(const char *name)
+int
+PyImport_ImportOrAddModule(const char *name, PyObject **pmodule)
 {
     PyObject *name_obj = PyUnicode_FromString(name);
     if (name_obj == NULL) {
-        return NULL;
+        *pmodule = NULL;
+        return -1;
     }
     PyThreadState *tstate = _PyThreadState_GET();
-    PyObject *module = import_add_module(tstate, name_obj);
+    int is_new;
+    PyObject *module = import_add_module(tstate, name_obj, &is_new);
     Py_DECREF(name_obj);
-    return module;
+
+    *pmodule = module;
+    if (module == NULL) {
+        assert(PyErr_Occurred());
+        return -1;
+    }
+    return is_new;
 }
 
 
 PyObject *
 PyImport_AddModuleObject(PyObject *name)
 {
     PyThreadState *tstate = _PyThreadState_GET();
-    PyObject *mod = import_add_module(tstate, name);
+    PyObject *mod = import_add_module(tstate, name, NULL);
     if (!mod) {
         return NULL;
     }
@@ -351,7 +367,7 @@ PyImport_AddModuleObject(PyObject *name)
     // 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.
+    // Use PyImport_ImportOrAddModule() to avoid these issues.
     PyObject *ref = PyWeakref_NewRef(mod, NULL);
     Py_DECREF(mod);
     if (ref == NULL) {
@@ -1258,7 +1274,7 @@ import_find_extension(PyThreadState *tstate, PyObject *name,
                 return NULL;
             }
         }
-        mod = import_add_module(tstate, name);
+        mod = import_add_module(tstate, name, NULL);
         if (mod == NULL) {
             return NULL;
         }
@@ -1386,7 +1402,7 @@ create_builtin(PyThreadState *tstate, PyObject *name, PyObject *spec)
         if (_PyUnicode_EqualToASCIIString(name, p->name)) {
             if (p->initfunc == NULL) {
                 /* Cannot re-init internal module ("sys" or "builtins") */
-                return import_add_module(tstate, name);
+                return import_add_module(tstate, name, NULL);
             }
             mod = (*p->initfunc)();
             if (mod == NULL) {
@@ -1671,7 +1687,7 @@ module_dict_for_exec(PyThreadState *tstate, PyObject *name)
 {
     PyObject *m, *d;
 
-    m = import_add_module(tstate, name);
+    m = import_add_module(tstate, name, NULL);
     if (m == NULL)
         return NULL;
     /* If the module is being reloaded, we get the old module back
@@ -2110,7 +2126,7 @@ PyImport_ImportFrozenModuleObject(PyObject *name)
     if (info.is_package) {
         /* Set __path__ to the empty list */
         PyObject *l;
-        m = import_add_module(tstate, name);
+        m = import_add_module(tstate, name, NULL);
         if (m == NULL)
             goto err_return;
         d = PyModule_GetDict(m);
@@ -2252,8 +2268,8 @@ init_importlib(PyThreadState *tstate, PyObject *sysmod)
         return -1;
     }
 
-    PyObject *importlib = PyImport_AddModuleRef("_frozen_importlib");
-    if (importlib == NULL) {
+    PyObject *importlib;
+    if (PyImport_ImportOrAddModule("_frozen_importlib", &importlib) < 0) {
         return -1;
     }
     IMPORTLIB(interp) = importlib;
@@ -3464,7 +3480,7 @@ _imp_init_frozen_impl(PyObject *module, PyObject *name)
     if (ret == 0) {
         Py_RETURN_NONE;
     }
-    return import_add_module(tstate, name);
+    return import_add_module(tstate, name, NULL);
 }
 
 /*[clinic input]
diff --git a/Python/pythonrun.c b/Python/pythonrun.c
@@ -277,8 +277,8 @@ PyRun_InteractiveOneObjectEx(FILE *fp, PyObject *filename,
         return parse_res;
     }
 
-    PyObject *main_module = PyImport_AddModuleRef("__main__");
-    if (main_module == NULL) {
+    PyObject *main_module;
+    if (PyImport_ImportOrAddModule("__main__", &main_module) < 0) {
         _PyArena_Free(arena);
         return -1;
     }
@@ -407,9 +407,10 @@ _PyRun_SimpleFileObject(FILE *fp, PyObject *filename, int closeit,
 {
     int ret = -1;
 
-    PyObject *main_module = PyImport_AddModuleRef("__main__");
-    if (main_module == NULL)
+    PyObject *main_module;
+    if (PyImport_ImportOrAddModule("__main__", &main_module) < 0) {
         return -1;
+    }
     PyObject *dict = PyModule_GetDict(main_module);  // borrowed ref
 
     int set_file_name = 0;
@@ -503,8 +504,8 @@ PyRun_SimpleFileExFlags(FILE *fp, const char *filename, int closeit,
 
 int
 _PyRun_SimpleStringFlagsWithName(const char *command, const char* name, PyCompilerFlags *flags) {
-    PyObject *main_module = PyImport_AddModuleRef("__main__");
-    if (main_module == NULL) {
+    PyObject *main_module;
+    if (PyImport_ImportOrAddModule("__main__", &main_module) < 0) {
         return -1;
     }
     PyObject *dict = PyModule_GetDict(main_module);  // borrowed ref

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	+Add :c:func:`PyImport_ImportOrAddModule` function and remove
	`2`	+:c:func:`PyImport_AddModuleRef` function. Patch by Victor Stinner.