gh-152298: Guard lazy resolve writebacks

Author: 1sgtpepper

Include/internal/pycore_lazyimportobject.h

@@ -33,6 +33,9 @@ PyAPI_FUNC(PyObject *) _PyLazyImport_GetName(PyObject *lazy_import);
 PyAPI_FUNC(PyObject *) _PyLazyImport_GetResolved(PyObject *lazy_import);
 PyAPI_FUNC(int) _PyLazyImport_FinishResolve(
     PyObject *lazy_import, PyObject *resolved);
+PyAPI_FUNC(int) _PyLazyImport_ReplaceDictItemIfCurrent(
+    PyObject *lazy_import, PyObject *dict, PyObject *name,
+    PyObject *resolved);
 PyAPI_FUNC(int) _PyLazyImport_SetGlobalBindingAndDictItem(
     PyObject *lazy_import, PyObject *globals, PyObject *name);
 PyAPI_FUNC(PyObject *) _PyLazyImport_New(

Lib/test/test_lazy_import/__init__.py

@@ -1244,6 +1244,127 @@ def test_resolve():
         self.assertEqual(result.returncode, 0, f"stdout: {result.stdout}, stderr: {result.stderr}")
         self.assertIn("OK", result.stdout)
 
+    def test_load_global_respects_rebound_global_during_reification(self):
+        """LOAD_GLOBAL should not overwrite a global rebound by the import hook."""
+        code = textwrap.dedent("""
+            import builtins
+            import types
+
+            real_import = builtins.__import__
+            calls = []
+            sentinel = object()
+
+            lazy import target_module as target
+
+            def custom_import(name, globals=None, locals=None, fromlist=None, level=0):
+                if name == "target_module":
+                    calls.append(name)
+                    globals["target"] = sentinel
+                    module = types.ModuleType(name)
+                    module.VALUE = "resolved"
+                    return module
+                return real_import(name, globals, locals, fromlist, level)
+
+            builtins.__import__ = custom_import
+            try:
+                def trigger_load_global():
+                    return target
+
+                resolved = trigger_load_global()
+                assert resolved.VALUE == "resolved"
+                assert globals()["target"] is sentinel
+                assert calls == ["target_module"], calls
+            finally:
+                builtins.__import__ = real_import
+            print("OK")
+        """)
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True
+        )
+        self.assertEqual(result.returncode, 0, f"stdout: {result.stdout}, stderr: {result.stderr}")
+        self.assertIn("OK", result.stdout)
+
+    def test_load_name_respects_deleted_global_during_reification(self):
+        """LOAD_NAME should not recreate a global deleted by the import hook."""
+        code = textwrap.dedent("""
+            import builtins
+            import types
+
+            real_import = builtins.__import__
+            calls = []
+
+            lazy import target_module as target
+
+            def custom_import(name, globals=None, locals=None, fromlist=None, level=0):
+                if name == "target_module":
+                    calls.append(name)
+                    del globals["target"]
+                    module = types.ModuleType(name)
+                    module.VALUE = "resolved"
+                    return module
+                return real_import(name, globals, locals, fromlist, level)
+
+            builtins.__import__ = custom_import
+            try:
+                resolved = target
+                assert resolved.VALUE == "resolved"
+                assert "target" not in globals()
+                assert calls == ["target_module"], calls
+            finally:
+                builtins.__import__ = real_import
+            print("OK")
+        """)
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True
+        )
+        self.assertEqual(result.returncode, 0, f"stdout: {result.stdout}, stderr: {result.stderr}")
+        self.assertIn("OK", result.stdout)
+
+    def test_module_attr_respects_rebound_global_during_reification(self):
+        """Module attribute access should not overwrite a rebound lazy binding."""
+        code = textwrap.dedent("""
+            import builtins
+            import types
+
+            real_import = builtins.__import__
+            calls = []
+            sentinel = object()
+            holder = types.ModuleType("holder")
+            holder.__dict__["__builtins__"] = builtins
+
+            exec("lazy import target_module as target", holder.__dict__)
+
+            def custom_import(name, globals=None, locals=None, fromlist=None, level=0):
+                if name == "target_module":
+                    calls.append(name)
+                    holder.__dict__["target"] = sentinel
+                    module = types.ModuleType(name)
+                    module.VALUE = "resolved"
+                    return module
+                return real_import(name, globals, locals, fromlist, level)
+
+            builtins.__import__ = custom_import
+            try:
+                resolved = holder.target
+                assert resolved.VALUE == "resolved"
+                assert holder.__dict__["target"] is sentinel
+                assert calls == ["target_module"], calls
+            finally:
+                builtins.__import__ = real_import
+            print("OK")
+        """)
+        result = subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True
+        )
+        self.assertEqual(result.returncode, 0, f"stdout: {result.stdout}, stderr: {result.stderr}")
+        self.assertIn("OK", result.stdout)
+
     def test_add_lazy_to_globals(self):
         code = textwrap.dedent("""
             import sys

Modules/_testinternalcapi/test_cases.c.h

@@ -184,6 +184,49 @@ _PyLazyImport_SetGlobalBindingAndDictItem(PyObject *op, PyObject *globals,
     return err;
 }
 
+static int
+lazy_import_replace_dict_item_if_current(PyObject *op, PyObject *globals,
+                                         PyObject *key, Py_hash_t key_hash,
+                                         PyObject *resolved)
+{
+    assert(PyLazyImport_CheckExact(op));
+    assert(PyDict_CheckExact(globals));
+    assert(!PyLazyImport_CheckExact(resolved));
+
+    PyObject *current = NULL;
+    int err = 0;
+
+    Py_BEGIN_CRITICAL_SECTION(globals);
+    int found = _PyDict_GetItemRef_KnownHash_LockHeld(
+        (PyDictObject *)globals, key, key_hash, &current);
+    if (found < 0) {
+        err = -1;
+    }
+    else if (found && current == op) {
+        err = _PyDict_SetItem_KnownHash_LockHeld(
+            (PyDictObject *)globals, key, resolved, key_hash);
+    }
+    Py_END_CRITICAL_SECTION();
+
+    Py_XDECREF(current);
+    return err;
+}
+
+int
+_PyLazyImport_ReplaceDictItemIfCurrent(PyObject *op, PyObject *globals,
+                                       PyObject *key, PyObject *resolved)
+{
+    assert(PyLazyImport_CheckExact(op));
+    assert(PyDict_CheckExact(globals));
+
+    Py_hash_t key_hash = PyObject_Hash(key);
+    if (key_hash == -1) {
+        return -1;
+    }
+    return lazy_import_replace_dict_item_if_current(
+        op, globals, key, key_hash, resolved);
+}
+
 int
 _PyLazyImport_FinishResolve(PyObject *op, PyObject *resolved)
 {
@@ -212,25 +255,12 @@ _PyLazyImport_FinishResolve(PyObject *op, PyObject *resolved)
         return 0;
     }
 
-    PyObject *current = NULL;
-
     if (globals != NULL) {
         assert(key != NULL);
         assert(PyDict_CheckExact(globals));
 
-        Py_BEGIN_CRITICAL_SECTION(globals);
-        int found = _PyDict_GetItemRef_KnownHash_LockHeld(
-            (PyDictObject *)globals, key, key_hash, &current);
-        if (found < 0) {
-            err = -1;
-        }
-        else if (found && current == op) {
-            err = _PyDict_SetItem_KnownHash_LockHeld(
-                (PyDictObject *)globals, key, resolved, key_hash);
-        }
-        Py_END_CRITICAL_SECTION();
-
-        Py_XDECREF(current);
+        err = lazy_import_replace_dict_item_if_current(
+            op, globals, key, key_hash, resolved);
         if (err < 0) {
             Py_DECREF(globals);
             Py_DECREF(key);

Objects/lazyimportobject.c

@@ -1368,7 +1368,9 @@ _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
                 return NULL;
             }
 
-            if (PyDict_SetItem(m->md_dict, name, new_value) < 0) {
+            if (_PyLazyImport_ReplaceDictItemIfCurrent(
+                    attr, m->md_dict, name, new_value) < 0)
+            {
                 Py_CLEAR(new_value);
             }
             Py_DECREF(attr);

Objects/moduleobject.c

@@ -2264,7 +2264,14 @@ dummy_func(
                     Py_DECREF(v_o);
                     ERROR_IF(true);
                 }
-                int err = PyDict_SetItem(GLOBALS(), name, l_v);
+                int err;
+                if (PyDict_CheckExact(GLOBALS())) {
+                    err = _PyLazyImport_ReplaceDictItemIfCurrent(
+                        v_o, GLOBALS(), name, l_v);
+                }
+                else {
+                    err = PyDict_SetItem(GLOBALS(), name, l_v);
+                }
                 if (err < 0) {
                     Py_DECREF(v_o);
                     Py_DECREF(l_v);

Python/bytecodes.c

@@ -3688,13 +3688,21 @@ _PyEval_LoadGlobalStackRef(PyObject *globals, PyObject *builtins, PyObject *name
     PyObject *res_o = PyStackRef_AsPyObjectBorrow(*writeto);
     if (res_o != NULL && PyLazyImport_CheckExact(res_o)) {
         PyObject *l_v = _PyImport_LoadLazyImportTstate(PyThreadState_GET(), res_o);
-        PyStackRef_CLOSE(writeto[0]);
         if (l_v == NULL) {
             assert(PyErr_Occurred());
+            PyStackRef_CLOSE(writeto[0]);
             *writeto = PyStackRef_NULL;
             return;
         }
-        int err = PyDict_SetItem(globals, name, l_v);
+        int err;
+        if (PyDict_CheckExact(globals)) {
+            err = _PyLazyImport_ReplaceDictItemIfCurrent(
+                res_o, globals, name, l_v);
+        }
+        else {
+            err = PyDict_SetItem(globals, name, l_v);
+        }
+        PyStackRef_CLOSE(writeto[0]);
         if (err < 0) {
             Py_DECREF(l_v);
             *writeto = PyStackRef_NULL;