Watch lazy import dicts during global specialization

Author: brandtbucher

Lib/test/test_lazy_import/__init__.py

@@ -949,6 +949,81 @@ def f():
         self.assertEqual(result.returncode, 0, f"stdout: {result.stdout}, stderr: {result.stderr}")
         self.assertIn("OK", result.stdout)
 
+    def test_add_lazy_to_exec_globals_after_specialization(self):
+        code = textwrap.dedent("""
+            source = '''
+            import sys
+            import types
+
+            lazy from test.test_lazy_import.data import basic2
+
+            assert 'test.test_lazy_import.data.basic2' not in sys.modules
+
+            class C: pass
+            sneaky = C()
+            sneaky.x = 1
+
+            def f():
+                t = 0
+                for _ in range(5):
+                    t += sneaky.x
+                return t
+
+            f()
+            globals()["sneaky"] = globals()["basic2"]
+            assert f() == 210
+            print("OK")
+            '''
+            ns = {"__name__": "lazy_exec_globals"}
+            exec(source, ns)
+        """)
+        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_exec_builtins_after_specialization(self):
+        code = textwrap.dedent("""
+            import builtins
+            source = '''
+            import sys
+            import types
+
+            lazy from test.test_lazy_import.data import basic2
+
+            assert 'test.test_lazy_import.data.basic2' not in sys.modules
+
+            class C: pass
+            sneaky = C()
+            sneaky.x = 1
+            __builtins__["sneaky"] = sneaky
+            del sneaky
+
+            def f():
+                t = 0
+                for _ in range(5):
+                    t += sneaky.x
+                return t
+
+            f()
+            __builtins__["sneaky"] = globals()["basic2"]
+            assert f() == 210
+            print("OK")
+            '''
+            ns = {"__name__": "lazy_exec_builtins", "__builtins__": builtins.__dict__.copy()}
+            exec(source, ns)
+        """)
+        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)
+
 
 @support.requires_subprocess()
 class MultipleNameFromImportTests(LazyImportTestCase):

Python/specialize.c

@@ -1386,6 +1386,7 @@ specialize_load_global_lock_held(
             SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
             goto fail;
         }
+        PyDict_Watch(MODULE_WATCHER_ID, globals);
 #ifdef Py_GIL_DISABLED
         maybe_enable_deferred_ref_count(value);
 #endif
@@ -1403,11 +1404,15 @@ specialize_load_global_lock_held(
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_LOAD_GLOBAL_NON_STRING_OR_SPLIT);
         goto fail;
     }
-    index = _PyDictKeys_StringLookup(builtin_keys, name);
+    index = _PyDict_LookupIndexAndValue((PyDictObject *)builtins, name, &value);
     if (index == DKIX_ERROR) {
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_EXPECTED_ERROR);
         goto fail;
     }
+    if (value != NULL && PyLazyImport_CheckExact(value)) {
+        SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_ATTR_MODULE_LAZY_VALUE);
+        goto fail;
+    }
     if (index != (uint16_t)index) {
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
         goto fail;
@@ -1422,6 +1427,7 @@ specialize_load_global_lock_held(
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
         goto fail;
     }
+    PyDict_Watch(MODULE_WATCHER_ID, globals);
     uint32_t builtins_version = _PyDict_GetKeysVersionForCurrentState(
             interp, (PyDictObject*) builtins);
     if (builtins_version == 0) {
@@ -1432,6 +1438,7 @@ specialize_load_global_lock_held(
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
         goto fail;
     }
+    PyDict_Watch(MODULE_WATCHER_ID, builtins);
     cache->index = (uint16_t)index;
     cache->module_keys_version = (uint16_t)globals_version;
     cache->builtin_keys_version = (uint16_t)builtins_version;