Use dictionary watchers to update version number on writes to dictionary

Author: DinoV

Include/internal/pycore_dict.h

@@ -265,6 +265,13 @@ static inline PyDictUnicodeEntry* DK_UNICODE_ENTRIES(PyDictKeysObject *dk) {
 #define DICT_UNIQUE_ID_SHIFT (32)
 #define DICT_UNIQUE_ID_MAX ((UINT64_C(1) << (64 - DICT_UNIQUE_ID_SHIFT)) - 1)
 
+/* The first three dict watcher IDs are reserved for CPython,
+ * so we don't need to check that they haven't been used */
+#define BUILTINS_WATCHER_ID     0
+#define GLOBALS_WATCHER_ID      1
+#define MODULE_WATCHER_ID       2
+#define FIRST_AVAILABLE_WATCHER 3
+
 
 PyAPI_FUNC(void)
 _PyDict_SendEvent(int watcher_bits,

Include/internal/pycore_moduleobject.h

@@ -19,6 +19,8 @@ extern int _PyModule_IsPossiblyShadowing(PyObject *);
 
 extern int _PyModule_IsExtension(PyObject *obj);
 
+extern int _PyModule_InitModuleDictWatcher(PyInterpreterState *interp);
+
 typedef int (*_Py_modexecfunc)(PyObject *);
 
 typedef struct {

Lib/test/test_import/test_lazy_imports.py

@@ -211,42 +211,9 @@ def test_lazy_import_type_exposed(self):
         self.assertHasAttr(types, 'LazyImportType')
         self.assertEqual(types.LazyImportType.__name__, 'lazy_import')
 
-    def test_lazy_import_type_invalid_name(self):
-        """passing invalid name to lazy imports should raise a TypeError"""
-        with self.assertRaises(TypeError) as cm:
-            types.LazyImportType({}, None)
-
-    def test_lazy_import_type_invalid_fromlist_type(self):
-        """LazyImportType should reject invalid fromlist types."""
-        # fromlist must be None, a string, or a tuple - not an int
-        with self.assertRaises(TypeError) as cm:
-            types.LazyImportType({}, "module", 0)
-        self.assertIn("fromlist must be None, a string, or a tuple", str(cm.exception))
-
-        # Also test with other invalid types
-        with self.assertRaises(TypeError):
-            types.LazyImportType({}, "module", [])  # list not allowed
-
-        with self.assertRaises(TypeError):
-            types.LazyImportType({}, "module", {"x": 1})  # dict not allowed
-
-    def test_lazy_import_type_valid_fromlist(self):
-        """LazyImportType should accept valid fromlist types."""
-        # None is valid (implicit)
-        lazy1 = types.LazyImportType({}, "module")
-        self.assertIsNotNone(lazy1)
-
-        # Explicit None is valid
-        lazy2 = types.LazyImportType({}, "module", None)
-        self.assertIsNotNone(lazy2)
-
-        # String is valid
-        lazy3 = types.LazyImportType({}, "module", "attr")
-        self.assertIsNotNone(lazy3)
-
-        # Tuple is valid
-        lazy4 = types.LazyImportType({}, "module", ("attr1", "attr2"))
-        self.assertIsNotNone(lazy4)
+    def test_lazy_import_type_cant_construct(self):
+        """LazyImportType should not be directly constructible."""
+        self.assertRaises(TypeError, types.LazyImportType, {}, "module")
 
 
 class SyntaxRestrictionTests(unittest.TestCase):
@@ -768,6 +735,38 @@ def test_resolve():
         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
+            import types
+
+            lazy from test.test_import.data.lazy_imports import basic2
+
+            assert 'test.test_import.data.lazy_imports.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")
+        """)
+        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)
+
 
 class MultipleNameFromImportTests(unittest.TestCase):
     """Tests for lazy from ... import with multiple names.

Objects/dictobject.c

@@ -7648,7 +7648,7 @@ PyDict_AddWatcher(PyDict_WatchCallback callback)
     PyInterpreterState *interp = _PyInterpreterState_GET();
 
     /* Start at 2, as 0 and 1 are reserved for CPython */
-    for (int i = 2; i < DICT_MAX_WATCHERS; i++) {
+    for (int i = FIRST_AVAILABLE_WATCHER; i < DICT_MAX_WATCHERS; i++) {
         if (!interp->dict_state.watchers[i]) {
             interp->dict_state.watchers[i] = callback;
             return i;

Objects/moduleobject.c

@@ -19,7 +19,6 @@
 
 #include "osdefs.h"               // MAXPATHLEN
 
-
 #define _PyModule_CAST(op) \
     (assert(PyModule_Check(op)), _Py_CAST(PyModuleObject*, (op)))
 
@@ -196,11 +195,40 @@ new_module_notrack(PyTypeObject *mt)
     return m;
 }
 
+/* Module dict watcher callback.
+ * When a module dictionary is modified, we need to clear the keys version
+ * to invalidate any cached lookups that depend on the dictionary structure.
+ */
+static int
+module_dict_watcher(PyDict_WatchEvent event, PyObject *dict,
+                    PyObject *key, PyObject *new_value)
+{
+    assert(PyDict_Check(dict));
+    // Only if a new lazy object shows up do we need to clear the dictionary. If
+    // this is adding a new key then the version will be reset anyway.
+    if (event == PyDict_EVENT_MODIFIED &&
+        new_value != NULL &&
+        PyLazyImport_CheckExact(new_value)) {
+        _PyDict_ClearKeysVersionLockHeld(dict);
+    }
+    return 0;
+}
+
+int
+_PyModule_InitModuleDictWatcher(PyInterpreterState *interp)
+{
+    // This is a reserved watcher for CPython so there's no need to check for non-NULL.
+    assert(interp->dict_state.watchers[MODULE_WATCHER_ID] == NULL);
+    interp->dict_state.watchers[MODULE_WATCHER_ID] = &module_dict_watcher;
+    return 0;
+}
+
 static void
 track_module(PyModuleObject *m)
 {
     _PyDict_EnablePerThreadRefcounting(m->md_dict);
     _PyObject_SetDeferredRefcount((PyObject *)m);
+    PyDict_Watch(MODULE_WATCHER_ID, m->md_dict);
     PyObject_GC_Track(m);
 }
 
@@ -1266,23 +1294,6 @@ _PyModule_IsPossiblyShadowing(PyObject *origin)
     return result;
 }
 
-int
-_PyModule_ReplaceLazyValue(PyObject *dict, PyObject *name, PyObject *value)
-{
-    // The adaptive interpreter uses the dictionary keys version to return the
-    // slot at a given index from the module. When replacing a value the
-    // version number doesn't change, so we need to atomically clear the
-    // version before replacing so that it doesn't return a lazy value.
-    int err;
-    Py_BEGIN_CRITICAL_SECTION(dict);
-
-    _PyDict_ClearKeysVersionLockHeld(dict);
-    err = _PyDict_SetItem_LockHeld((PyDictObject *)dict, name, value);
-
-    Py_END_CRITICAL_SECTION();
-    return err;
-}
-
 PyObject*
 _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
 {
@@ -1306,7 +1317,7 @@ _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
                 return NULL;
             }
 
-            if (_PyModule_ReplaceLazyValue(m->md_dict, name, new_value) < 0) {
+            if (PyDict_SetItem(m->md_dict, name, new_value) < 0) {
                 Py_CLEAR(new_value);
             }
             Py_DECREF(attr);

Python/bytecodes.c

@@ -1829,7 +1829,7 @@ dummy_func(
                     Py_DECREF(v_o);
                     ERROR_IF(true);
                 }
-                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                int err = PyDict_SetItem(GLOBALS(), name, l_v);
                 if (err < 0) {
                     Py_DECREF(v_o);
                     Py_DECREF(l_v);

Python/ceval.c

@@ -4186,7 +4186,7 @@ _PyEval_LoadGlobalStackRef(PyObject *globals, PyObject *builtins, PyObject *name
             *writeto = PyStackRef_NULL;
             return;
         }
-        int err = _PyModule_ReplaceLazyValue(globals, name, l_v);
+        int err = PyDict_SetItem(globals, name, l_v);
         if (err < 0) {
             Py_DECREF(l_v);
             *writeto = PyStackRef_NULL;

Python/executor_cases.c.h

@@ -5530,8 +5530,8 @@ publish_lazy_imports_on_module(PyThreadState *tstate,
         }
 
         // Publish on the module that was just imported.
-        if (_PyModule_ReplaceLazyValue(module_dict, attr_name,
-                                       lazy_module_attr) < 0) {
+        if (PyDict_SetItem(module_dict, attr_name,
+                           lazy_module_attr) < 0) {
             Py_DECREF(lazy_module_attr);
             Py_DECREF(attr_name);
             return -1;