[mypyc] Fix missing module during concurrent circular import

mypyc/irbuild/builder.py

@@ -523,7 +523,8 @@ def add_to_non_ext_dict(
         # doesn't cause contention.
         self.builder.set_immortal_if_free_threaded(val, line)
 
-    def gen_import(self, module: str, line: int) -> None:
+    def gen_import(self, module: str, line: int) -> Value | None:
+        """Import a module and return it for a same-group native import."""
         self.imports[module] = None
 
         needs_import, out = BasicBlock(), BasicBlock()
@@ -547,7 +548,20 @@ def gen_import(self, module: str, line: int) -> None:
                 module_lock_api = Integer(0, c_pointer_rprimitive)
         else:
             module_id = self.load_str(module, line)
-        self.check_if_module_loaded(module_id, module_cache, line, needs_import, out, import_state)
+        already_loaded = BasicBlock() if is_same_group_native else out
+        self.check_if_module_loaded(
+            module_id, module_cache, line, needs_import, already_loaded, import_state
+        )
+
+        result: Register | None = None
+        if is_same_group_native:
+            # Preserve the module returned by the native import. In particular,
+            # don't look it up again in sys.modules after a circular import:
+            # importlib temporarily removes a module there during finalization.
+            result = Register(object_rprimitive, line=line)
+            self.activate_block(already_loaded)
+            self.add(Assign(result, self.get_module(module, line), line))
+            self.goto(out)
 
         self.activate_block(needs_import)
         if is_same_group_native:
@@ -593,13 +607,16 @@ def gen_import(self, module: str, line: int) -> None:
                 ],
                 line,
             )
+            assert result is not None
+            self.add(Assign(result, value, line))
         else:
             # Import using generic Python C API
             assert module_id is not None
             value = self.call_c(import_op, [module_id], line)
         if not is_same_group_native:
             self.call_c(import_cache_replace_for_import_op, [module_cache, value], line)
         self.goto_and_activate(out)
+        return result
 
     def check_if_module_loaded(
         self,
@@ -643,8 +660,8 @@ def get_module_attr(self, module: str, attr: str, line: int) -> Value:
 
         Import the module if needed.
         """
-        self.gen_import(module, line)
-        module_obj = self.get_module(module, line)
+        imported = self.gen_import(module, line)
+        module_obj = imported if imported is not None else self.get_module(module, line)
         return self.py_get_attr(module_obj, attr, line)
 
     def assign_if_null(self, target: Register, get_val: Callable[[], Value], line: int) -> None:

mypyc/irbuild/statement.py

@@ -419,14 +419,16 @@ def transform_imports_without_grouping(
 ) -> None:
     globals = builder.load_globals_dict()
     for mod_id, as_name, line in group:
-        builder.gen_import(mod_id, line)
+        imported = builder.gen_import(mod_id, line)
         globals_id, globals_name = import_globals_id_and_name(mod_id, as_name)
+        globals_name_value = builder.load_str(globals_name)
+        module = (
+            imported
+            if imported is not None and globals_id == mod_id
+            else builder.get_module(globals_id, line)
+        )
         builder.gen_method_call(
-            globals,
-            "__setitem__",
-            [builder.load_str(globals_name), builder.get_module(globals_id, line)],
-            result_type=None,
-            line=line,
+            globals, "__setitem__", [globals_name_value, module], result_type=None, line=line
         )
 
 
@@ -558,7 +560,8 @@ def transform_import_from_buckets(
             ]
             transform_imports_without_grouping(builder, group)
         elif bucket.kind == IMPORT_NATIVE_ATTR:
-            builder.gen_import(module_id, line)
+            native_module = builder.gen_import(module_id, line)
+            assert native_module is not None
             names_literal = builder.add(LoadLiteral(tuple(bucket.names), object_rprimitive))
             if bucket.as_names == bucket.names:
                 as_names_literal = names_literal
@@ -568,12 +571,7 @@ def transform_import_from_buckets(
                 )
             builder.call_c(
                 get_native_attrs_op,
-                [
-                    builder.load_str(module_id),
-                    names_literal,
-                    as_names_literal,
-                    builder.load_globals_dict(),
-                ],
+                [native_module, names_literal, as_names_literal, builder.load_globals_dict()],
                 line,
             )
         else:

mypyc/lib-rt/CPy.h

@@ -1085,7 +1085,7 @@ bool CPyImport_ImportMany(PyObject *modules, CPyModuleCache *statics[], PyObject
                           PyObject *tb_path, PyObject *tb_function, Py_ssize_t *tb_lines);
 PyObject *CPyImport_ImportFromMany(PyObject *mod_id, PyObject *names, PyObject *as_names,
                                    PyObject *globals);
-PyObject *CPyImport_GetNativeAttrs(PyObject *mod_id, PyObject *names, PyObject *as_names,
+PyObject *CPyImport_GetNativeAttrs(PyObject *mod, PyObject *names, PyObject *as_names,
                                    PyObject *globals);
 PyObject *CPyImport_ImportNative(PyObject *module_name,
                                  PyObject *(*init_only_fn)(void),

mypyc/lib-rt/misc_ops.c

@@ -906,33 +906,24 @@ PyObject *CPyImport_ImportFromMany(PyObject *mod_id, PyObject *names, PyObject *
     return mod;
 }
 
-// Import attributes from an already-imported native module and store them
+// Import attributes from an already-imported native module object and store them
 // in the globals dict.  Returns the module on success, NULL on error.
-PyObject *CPyImport_GetNativeAttrs(PyObject *mod_id, PyObject *names,
+PyObject *CPyImport_GetNativeAttrs(PyObject *mod, PyObject *names,
                                    PyObject *as_names, PyObject *globals) {
-    PyObject *mod = PyImport_GetModule(mod_id);
-    if (mod == NULL) {
-        if (!PyErr_Occurred()) {
-            PyErr_Format(PyExc_ImportError, "module '%U' is not in sys.modules", mod_id);
-        }
-        return NULL;
-    }
     for (Py_ssize_t i = 0; i < PyTuple_GET_SIZE(names); i++) {
         PyObject *name = PyTuple_GET_ITEM(names, i);
         PyObject *as_name = PyTuple_GET_ITEM(as_names, i);
         PyObject *obj = PyObject_GetAttr(mod, name);
         if (obj == NULL) {
-            Py_DECREF(mod);
             return NULL;
         }
         int ret = CPyDict_SetItem(globals, as_name, obj);
         Py_DECREF(obj);
         if (ret < 0) {
-            Py_DECREF(mod);
             return NULL;
         }
     }
-    return mod;
+    return Py_NewRef(mod);
 }
 
 // From CPython
@@ -1667,6 +1658,18 @@ static int CPyImport_ReleaseLockPreservingException(PyObject *module_lock) {
     return result;
 }
 
+static void CPyImport_SetModuleReplacedError(PyObject *module_name) {
+    PyErr_Format(PyExc_ImportError,
+                 "native module '%U' in sys.modules was replaced after initialization",
+                 module_name);
+}
+
+static void CPyImport_SetModuleMissingError(PyObject *module_name) {
+    PyErr_Format(PyExc_ImportError,
+                 "initialized native module '%U' is missing from sys.modules",
+                 module_name);
+}
+
 // Execute a module once; caller holds the module lock.
 int CPyImport_Exec(PyObject *module, int (*exec_fn)(PyObject *), CPyImportState *state,
                    CPyModuleCache *module_cache) {
@@ -1721,15 +1724,27 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
     // We then check sys.modules to determine whether the module body
     // has already been executed (or is being executed in a circular import).
     PyObject *module_lock;
-    int lock_result = CPyImport_AcquireLock(lock_api, module_name, &module_lock);
-    if (lock_result == CPY_LOCK_DEADLOCK) {
+    int lock_result;
+    bool retried_after_execution = false;
+    while ((lock_result = CPyImport_AcquireLock(lock_api, module_name, &module_lock)) ==
+           CPY_LOCK_DEADLOCK) {
         PyObject *partial = PyDict_GetItemWithError(PyImport_GetModuleDict(), module_name);
         if (partial != NULL &&
                 (*module_static == NULL || partial == (PyObject *)*module_static)) {
             CPyImport_ReplaceModuleCacheUnverified(module_cache, partial);
             Py_INCREF(partial);
             return partial;
         }
+        // Multiple threads can detect the same cycle concurrently on a
+        // free-threaded build. If the target body has finished but the module
+        // is temporarily absent during importlib's pop/reinsert, retry once to
+        // wait for finalization. If it's still missing quit and report an error
+        // to not spin infinitely.
+        if (partial == NULL && !PyErr_Occurred() && CPyImport_IsExecuted(state) &&
+                !retried_after_execution) {
+            retried_after_execution = true;
+            continue;
+        }
         if (!PyErr_Occurred()) {
             PyErr_Format(PyExc_ImportError,
                          "import deadlock for native module '%U' without a partial module",
@@ -1779,9 +1794,7 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
                 }
                 return existing;
             }
-            PyErr_Format(PyExc_ImportError,
-                         "native module '%U' in sys.modules was replaced after initialization",
-                         module_name);
+            CPyImport_SetModuleReplacedError(module_name);
             CPyImport_ReleaseLockPreservingException(module_lock);
             return NULL;
         }
@@ -1793,9 +1806,7 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
 
     if (CPyImport_IsInitialized(state) ||
             CPyImport_GetModuleCache(module_cache) == (PyObject *)*module_static) {
-        PyErr_Format(PyExc_ImportError,
-                     "initialized native module '%U' is missing from sys.modules",
-                     module_name);
+        CPyImport_SetModuleMissingError(module_name);
         CPyImport_ReleaseLockPreservingException(module_lock);
         return NULL;
     }
@@ -1831,6 +1842,20 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
         goto fail;
     }
 
+    // Direct imports bypass importlib's final sys.modules lookup. Reject changes
+    // to the entry instead of returning an object different from future imports.
+    PyObject *current = PyDict_GetItemWithError(module_dict, module_name);
+    if (current == NULL) {
+        if (!PyErr_Occurred()) {
+            CPyImport_SetModuleMissingError(module_name);
+        }
+        goto fail;
+    }
+    if (current != modobj) {
+        CPyImport_SetModuleReplacedError(module_name);
+        goto fail;
+    }
+
     end_result = CPyImport_EndInitializing(initializing_spec);
     initializing_spec = NULL;
     if (end_result < 0) {

mypyc/primitives/misc_ops.py

@@ -207,7 +207,7 @@
     error_kind=ERR_MAGIC,
 )
 
-# Get attributes from an already-imported native module and store them in globals.
+# Get attributes from an already-imported native module object and store them in globals.
 get_native_attrs_op = custom_op(
     arg_types=[object_rprimitive, object_rprimitive, object_rprimitive, object_rprimitive],
     return_type=object_rprimitive,

mypyc/test-data/run-multimodule.test

@@ -438,6 +438,63 @@ else:
 
 assert sys.modules["other_top"] is fake
 
+[case testNativeImportRejectsSysModulesReplacementDuringInitialization]
+# separate: [(["native.py", "other_target.py"], "testgroup")]
+def load_target() -> None:
+    import other_target
+
+[file other_target.py]
+import sys
+
+replacement = object()
+sys.modules[__name__] = replacement
+
+[file sys/__init__.pyi]
+modules: dict[str, object]
+
+[file driver.py]
+import sys
+
+import native
+
+from testutil import assertRaises
+
+with assertRaises(
+    ImportError,
+    "native module 'other_target' in sys.modules was replaced after initialization",
+):
+    native.load_target()
+
+assert "other_target" not in sys.modules
+
+[case testNativeImportRejectsSysModulesRemovalDuringInitialization]
+# separate: [(["native.py", "other_target.py"], "testgroup")]
+def load_target() -> None:
+    import other_target
+
+[file other_target.py]
+import sys
+
+del sys.modules[__name__]
+
+[file sys/__init__.pyi]
+modules: dict[str, object]
+
+[file driver.py]
+import sys
+
+import native
+
+from testutil import assertRaises
+
+with assertRaises(
+    ImportError,
+    "initialized native module 'other_target' is missing from sys.modules",
+):
+    native.load_target()
+
+assert "other_target" not in sys.modules
+
 [case testMultiModuleFastpaths]
 [file other_main.py]
 
@@ -2633,6 +2690,48 @@ with ThreadPoolExecutor(max_workers=2) as executor:
 assert other_a.other_value == "b"
 assert other_b.other_value == "a"
 
+[case testConcurrentCircularNativeFromImports]
+# separate: [(["other_a.py", "other_b.py"], "testgroup")]
+pass
+
+[file import_sync.py]
+from threading import Barrier
+
+barrier = Barrier(2)
+
+[file other_a.py]
+from import_sync import barrier
+
+def value() -> str:
+    return "a"
+
+barrier.wait(timeout=15)
+from other_b import value as other_value_fn
+other_value = other_value_fn()
+
+[file other_b.py]
+from import_sync import barrier
+
+def value() -> str:
+    return "b"
+
+barrier.wait(timeout=15)
+from other_a import value as other_value_fn
+other_value = other_value_fn()
+
+[file driver.py]
+from concurrent.futures import ThreadPoolExecutor
+import importlib
+
+with ThreadPoolExecutor(max_workers=2) as executor:
+    future_a = executor.submit(importlib.import_module, "other_a")
+    future_b = executor.submit(importlib.import_module, "other_b")
+    other_a = future_a.result(timeout=15)
+    other_b = future_b.result(timeout=15)
+
+assert other_a.other_value == "b"
+assert other_b.other_value == "a"
+
 [case testNativeModuleReimportBehavior]
 value = 42