Ensure object in sys.modules is the native module

Author: JukkaL

mypyc/irbuild/builder.py

@@ -62,6 +62,7 @@
     BITMAP_BITS,
     EXT_SUFFIX,
     GENERATOR_ATTRIBUTE_PREFIX,
+    MODULE_PREFIX,
     SELF_NAME,
     TEMP_ATTR_NAME,
     shared_lib_name,
@@ -83,6 +84,7 @@
     InitStatic,
     Integer,
     IntOp,
+    LoadAddress,
     LoadGlobal,
     LoadStatic,
     MethodCall,
@@ -115,6 +117,7 @@
     is_tagged,
     is_tuple_rprimitive,
     none_rprimitive,
+    object_pointer_rprimitive,
     object_rprimitive,
     str_rprimitive,
 )
@@ -491,6 +494,12 @@ def gen_import(self, module: str, line: int) -> None:
             exec_func = self.add(
                 LoadGlobal(c_pointer_rprimitive, f"CPyExec_{exported_name(module)}")
             )
+            module_static = self.add(
+                LoadAddress(
+                    object_pointer_rprimitive,
+                    f"{MODULE_PREFIX}{exported_name(module + '__internal')}",
+                )
+            )
             group_name = self.mapper.group_map.get(self.module_name)
             if group_name is not None:
                 shared_lib_mod_name = shared_lib_name(group_name)
@@ -509,6 +518,7 @@ def gen_import(self, module: str, line: int) -> None:
                     self.load_str(module, line),
                     init_only_func,
                     exec_func,
+                    module_static,
                     shared_lib_file,
                     ext_suffix,
                     Integer(1 if is_pkg else 0, c_pyssize_t_rprimitive),

mypyc/lib-rt/CPy.h

@@ -964,6 +964,7 @@ PyObject *CPyImport_GetNativeAttrs(PyObject *mod_id, PyObject *names, PyObject *
 PyObject *CPyImport_ImportNative(PyObject *module_name,
                                  PyObject *(*init_only_fn)(void),
                                  int (*exec_fn)(PyObject *),
+                                 CPyModule **module_static,
                                  PyObject *shared_lib_file, PyObject *ext_suffix,
                                  Py_ssize_t is_package);
 

mypyc/lib-rt/misc_ops.c

@@ -1410,10 +1410,12 @@ static int CPyImport_SetModuleSpec(PyObject *modobj, PyObject *module_name,
 PyObject *CPyImport_ImportNative(PyObject *module_name,
                                  PyObject *(*init_only_fn)(void),
                                  int (*exec_fn)(PyObject *),
+                                 CPyModule **module_static,
                                  PyObject *shared_lib_file, PyObject *ext_suffix,
                                  Py_ssize_t is_package) {
     PyObject *parent_module = NULL;
     PyObject *child_name = NULL;
+    PyObject *exc_type, *exc_val, *exc_tb;
     Py_ssize_t name_len = PyUnicode_GetLength(module_name);
     if (name_len < 0) {
         return NULL;
@@ -1449,27 +1451,45 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
         return NULL;
     }
 
-    PyObject *modobj = init_only_fn();
-    if (modobj == NULL) {
+    PyObject *existing = PyDict_GetItemWithError(module_dict, module_name);
+    if (existing != NULL) {
+        if (*module_static != NULL) {
+            if (existing == (PyObject *)*module_static) {
+                Py_INCREF(existing);
+                Py_XDECREF(parent_module);
+                Py_XDECREF(child_name);
+                return existing;
+            }
+            PyErr_Format(PyExc_ImportError,
+                         "native module '%U' in sys.modules was replaced after initialization",
+                         module_name);
+            Py_XDECREF(parent_module);
+            Py_XDECREF(child_name);
+            return NULL;
+        }
+    }
+    if (PyErr_Occurred()) {
         Py_XDECREF(parent_module);
         Py_XDECREF(child_name);
         return NULL;
     }
 
-    // Check if the module was already imported (e.g. via CPyInit_* from the
-    // standard import machinery, or by a previous CPyImport_ImportNative call).
-    int already_imported = PyDict_Contains(module_dict, module_name);
-    if (already_imported < 0) {
+    PyObject *modobj = init_only_fn();
+    if (modobj == NULL) {
         Py_XDECREF(parent_module);
         Py_XDECREF(child_name);
-        Py_DECREF(modobj);
         return NULL;
     }
 
-    if (!already_imported) {
-        if (PyObject_SetItem(module_dict, module_name, modobj) < 0) {
-            goto fail;
-        }
+    if (PyObject_SetItem(module_dict, module_name, modobj) < 0) {
+        goto fail;
+    }
+
+    if (*module_static != (CPyModule *)modobj) {
+        PyErr_Format(PyExc_ImportError,
+                     "native module '%U' was initialized inconsistently",
+                     module_name);
+        goto fail;
     }
 
     // Set __package__ before executing the module body so it is available
@@ -1512,15 +1532,6 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
         goto fail;
     }
 
-    // If the module was already imported, skip executing the module body
-    // (it was already executed). This handles circular imports and modules
-    // first imported via the standard Python import machinery (CPyInit_*).
-    if (already_imported) {
-        Py_XDECREF(parent_module);
-        Py_XDECREF(child_name);
-        return modobj;
-    }
-
     // Now execute the module body, with __file__ and __package__ already set.
     if (exec_fn(modobj) != 0) {
         goto fail;
@@ -1537,15 +1548,12 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
     return modobj;
 
 fail:
-    // Clean up on failure: if we added the module to sys.modules, remove it
-    // so that a subsequent import attempt will retry initialization.
-    if (!already_imported) {
-        PyObject *exc_type, *exc_val, *exc_tb;
-        PyErr_Fetch(&exc_type, &exc_val, &exc_tb);
-        PyObject_DelItem(module_dict, module_name);
-        PyErr_Clear();
-        PyErr_Restore(exc_type, exc_val, exc_tb);
-    }
+    // Clean up on failure so that a subsequent import attempt will retry
+    // initialization.
+    PyErr_Fetch(&exc_type, &exc_val, &exc_tb);
+    PyObject_DelItem(module_dict, module_name);
+    PyErr_Clear();
+    PyErr_Restore(exc_type, exc_val, exc_tb);
     Py_XDECREF(parent_module);
     Py_XDECREF(child_name);
     Py_DECREF(modobj);

mypyc/primitives/misc_ops.py

@@ -139,11 +139,13 @@
 
 # Import a native same-group module directly via C-level init/exec functions.
 native_import_op = custom_op(
-    # (module name, init-only function, exec function, shared lib __file__, ext suffix, is_package)
+    # (module name, init-only function, exec function, module static,
+    #  shared lib __file__, ext suffix, is_package)
     arg_types=[
         str_rprimitive,
         c_pointer_rprimitive,
         c_pointer_rprimitive,
+        object_pointer_rprimitive,
         object_rprimitive,
         str_rprimitive,
         c_pyssize_t_rprimitive,

mypyc/test-data/run-multimodule.test

@@ -369,6 +369,49 @@ else:
 assert not hasattr(other_pkg, "other_sub")
 assert "other_pkg.other_sub" not in sys.modules
 
+[case testNativeImportUsesExistingNativeSysModulesObject]
+# separate: [(["native.py", "other_top.py"], "testgroup")]
+import other_top
+
+def f() -> int:
+    return other_top.value
+[file other_top.py]
+value = 42
+[file driver.py]
+import native
+import other_top
+
+assert native.other_top is other_top
+assert native.f() == 42
+
+[case testNativeImportRejectsForeignSysModulesObject]
+# separate: [(["native.py", "other_top.py"], "testgroup")]
+import other_top
+
+def f() -> int:
+    return other_top.value
+[file other_top.py]
+value = 42
+[file driver.py]
+import sys
+import types
+
+import other_top
+
+fake = types.ModuleType("other_top")
+fake.value = -1
+sys.modules["other_top"] = fake
+
+try:
+    import native
+except ImportError as e:
+    assert "other_top" in str(e)
+    assert "replaced after initialization" in str(e)
+else:
+    assert False, "import native should fail"
+
+assert sys.modules["other_top"] is fake
+
 [case testMultiModuleFastpaths]
 [file other_main.py]