Merge pull request python#9 from LazyImportsCabal/cleanup

pablogsal · web-flow · commit fe526b4e548c · 2025-10-09T00:24:09.000+01:00
More PEP-810 compatibility, fix memory leaks and LOAD_ATTR specialization
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -38,6 +38,7 @@ extern PyObject* _PyDict_GetItemIdWithError(PyObject *dp,
 extern int _PyDict_ContainsId(PyObject *, _Py_Identifier *);
 extern int _PyDict_SetItemId(PyObject *dp, _Py_Identifier *key, PyObject *item);
 extern int _PyDict_DelItemId(PyObject *mp, _Py_Identifier *key);
+extern void _PyDict_ClearKeysVersion(PyObject *mp);
 
 extern int _PyDict_Next(
     PyObject *mp, Py_ssize_t *pos, PyObject **key, PyObject **value, Py_hash_t *hash);
diff --git a/Include/internal/pycore_lazyimportobject.h b/Include/internal/pycore_lazyimportobject.h
@@ -13,7 +13,7 @@ PyAPI_DATA(PyTypeObject) PyLazyImport_Type;
 
 typedef struct {
     PyObject_HEAD
-    PyObject *lz_import_func;
+    PyObject *lz_builtins;
     PyObject *lz_from;
     PyObject *lz_attr;
     /* Frame information for the original import location */
diff --git a/Include/internal/pycore_moduleobject.h b/Include/internal/pycore_moduleobject.h
@@ -58,6 +58,8 @@ extern Py_ssize_t _PyModule_GetFilenameUTF8(
 PyObject* _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress);
 PyObject* _Py_module_getattro(PyObject *m, PyObject *name);
 
+PyAPI_FUNC(int) _PyModule_ReplaceLazyValue(PyObject *dict, PyObject *name, PyObject *value);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/Lib/importlib/__init__.py b/Lib/importlib/__init__.py
@@ -59,7 +59,7 @@
 
 from ._bootstrap import __import__
 
-from _imp import lazy_import, set_lazy_imports
+from _imp import set_lazy_imports
 
 
 def invalidate_caches():
diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
@@ -2790,7 +2790,7 @@ def test_lazy_import_pkg_cross_import(self):
 
         g = test.test_import.data.lazy_imports.pkg.c.get_globals()
         self.assertEqual(type(g["x"]), int)
-        self.assertEqual(type(g["b"]), importlib.lazy_import)
+        self.assertEqual(type(g["b"]), types.LazyImportType)
 
 class TestSinglePhaseSnapshot(ModuleSnapshot):
     """A representation of a single-phase init module for testing.
diff --git a/Modules/_typesmodule.c b/Modules/_typesmodule.c
@@ -2,6 +2,7 @@
 
 #include "Python.h"
 #include "pycore_descrobject.h"   // _PyMethodWrapper_Type
+#include "pycore_lazyimportobject.h" // PyLazyImport_Type
 #include "pycore_namespace.h"     // _PyNamespace_Type
 #include "pycore_object.h"        // _PyNone_Type, _PyNotImplemented_Type
 #include "pycore_unionobject.h"   // _PyUnion_Type
@@ -35,6 +36,7 @@ _types_exec(PyObject *m)
     EXPORT_STATIC_TYPE("GetSetDescriptorType", PyGetSetDescr_Type);
     // LambdaType is the same as FunctionType
     EXPORT_STATIC_TYPE("LambdaType", PyFunction_Type);
+    EXPORT_STATIC_TYPE("LazyImportType", PyLazyImport_Type);
     EXPORT_STATIC_TYPE("MappingProxyType", PyDictProxy_Type);
     EXPORT_STATIC_TYPE("MemberDescriptorType", PyMemberDescr_Type);
     EXPORT_STATIC_TYPE("MethodDescriptorType", PyMethodDescr_Type);
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -4719,6 +4719,14 @@ sizeof_lock_held(PyDictObject *mp)
     return (Py_ssize_t)res;
 }
 
+void
+_PyDict_ClearKeysVersion(PyObject *mp)
+{
+    ASSERT_DICT_LOCKED(mp);
+
+    FT_ATOMIC_STORE_UINT32_RELAXED(((PyDictObject *)mp)->ma_keys->dk_version, 0);
+}
+
 Py_ssize_t
 _PyDict_SizeOf(PyDictObject *mp)
 {
diff --git a/Objects/lazyimportobject.c b/Objects/lazyimportobject.c
@@ -8,7 +8,7 @@
 #include "pycore_interpframe.h"
 
 PyObject *
-_PyLazyImport_New(PyObject *import_func, PyObject *from, PyObject *attr)
+_PyLazyImport_New(PyObject *builtins, PyObject *from, PyObject *attr)
 {
     PyLazyImportObject *m;
     if (!from || !PyUnicode_Check(from)) {
@@ -23,8 +23,8 @@ _PyLazyImport_New(PyObject *import_func, PyObject *from, PyObject *attr)
     if (m == NULL) {
         return NULL;
     }
-    Py_XINCREF(import_func);
-    m->lz_import_func = import_func;
+    Py_XINCREF(builtins);
+    m->lz_builtins = builtins;
     Py_INCREF(from);
     m->lz_from = from;
     Py_XINCREF(attr);
@@ -52,7 +52,7 @@ static void
 lazy_import_dealloc(PyLazyImportObject *m)
 {
     PyObject_GC_UnTrack(m);
-    Py_XDECREF(m->lz_import_func);
+    Py_XDECREF(m->lz_builtins);
     Py_XDECREF(m->lz_from);
     Py_XDECREF(m->lz_attr);
     Py_XDECREF(m->lz_code);
@@ -88,7 +88,7 @@ lazy_import_repr(PyLazyImportObject *m)
 static int
 lazy_import_traverse(PyLazyImportObject *m, visitproc visit, void *arg)
 {
-    Py_VISIT(m->lz_import_func);
+    Py_VISIT(m->lz_builtins);
     Py_VISIT(m->lz_from);
     Py_VISIT(m->lz_attr);
     Py_VISIT(m->lz_code);
@@ -98,7 +98,7 @@ lazy_import_traverse(PyLazyImportObject *m, visitproc visit, void *arg)
 static int
 lazy_import_clear(PyLazyImportObject *m)
 {
-    Py_CLEAR(m->lz_import_func);
+    Py_CLEAR(m->lz_builtins);
     Py_CLEAR(m->lz_from);
     Py_CLEAR(m->lz_attr);
     Py_CLEAR(m->lz_code);
@@ -144,8 +144,8 @@ static PyMethodDef lazy_methods[] = {
 
 PyTypeObject PyLazyImport_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    "lazy_import",                              /* tp_name */
-    sizeof(PyLazyImportObject),                       /* tp_basicsize */
+    "LazyImport",                               /* tp_name */
+    sizeof(PyLazyImportObject),                 /* tp_basicsize */
     0,                                          /* tp_itemsize */
     (destructor)lazy_import_dealloc,            /* tp_dealloc */
     0,                                          /* tp_print */
diff --git a/Objects/moduleobject.c b/Objects/moduleobject.c
@@ -1036,6 +1036,23 @@ _PyModule_IsPossiblyShadowing(PyObject *origin)
     return result;
 }
 
+int
+_PyModule_ReplaceLazyValue(PyObject *dict, PyObject *name, PyObject *value)
+{
+    // The adaptive interpreter uses the dictionary 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_ClearKeysVersion(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)
 {
@@ -1052,13 +1069,13 @@ _Py_module_getattro_impl(PyModuleObject *m, PyObject *name, int suppress)
                     // instead treated as if the attribute doesn't exist.
                     PyErr_Clear();
                 }
-
-                return NULL;
-            } else if (PyDict_SetItem(m->md_dict, name, new_value) < 0) {
-                Py_DECREF(new_value);
                 Py_DECREF(attr);
                 return NULL;
             }
+
+            if (_PyModule_ReplaceLazyValue(m->md_dict, name, new_value) < 0) {
+                Py_CLEAR(new_value);
+            }
             Py_DECREF(attr);
             return new_value;
         }
@@ -1490,7 +1507,7 @@ module_get_dict(PyObject *mod, void *Py_UNUSED(ignored))
                     if (new_value == NULL) {
                         return NULL;
                     }
-                    if (PyDict_SetItem(self->md_dict, key, new_value) < 0) {
+                    if (_PyModule_ReplaceLazyValue(self->md_dict, key, new_value) < 0) {
                         Py_DECREF(new_value);
                         return NULL;
                     }
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
@@ -1799,12 +1799,13 @@ dummy_func(
             ERROR_IF(v_o == NULL);
             if (PyLazyImport_CheckExact(v_o)) {
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, v_o);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
+                Py_DECREF(v_o);
+                ERROR_IF(l_v == NULL);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                if (err < 0) {
                     JUMP_TO_LABEL(error);
                 }
-                Py_DECREF(v_o);
                 v_o = l_v;
-                ERROR_IF(v_o == NULL);
             }
 
             v = PyStackRef_FromPyObjectSteal(v_o);
@@ -1838,13 +1839,14 @@ dummy_func(
             PyObject *res_o = PyStackRef_AsPyObjectBorrow(*res);
             if (PyLazyImport_CheckExact(res_o)) {
                 PyObject *l_v = _PyImport_LoadLazyImportTstate(tstate, res_o);
-                if (l_v != NULL && PyDict_SetItem(GLOBALS(), name, l_v) < 0) {
+                PyStackRef_CLOSE(res[0]);
+                ERROR_IF(l_v == NULL);
+                int err = _PyModule_ReplaceLazyValue(GLOBALS(), name, l_v);
+                if (err < 0) {
+                    Py_DECREF(l_v);
                     JUMP_TO_LABEL(error);
                 }
-                res_o = l_v;
-                PyStackRef_CLOSE(res[0]);
-                ERROR_IF(res_o == NULL);
-                *res = PyStackRef_FromPyObjectSteal(res_o);
+                *res = PyStackRef_FromPyObjectSteal(l_v);
             }
         }
 
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -3038,6 +3038,9 @@ check_lazy_import_comatibility(PyThreadState *tstate, PyObject *globals,
      // Check if this module should be imported lazily due to the compatbility mode support via
     // __lazy_modules__.
     PyObject *lazy_modules = NULL;
+    PyObject *abs_name = NULL;
+    int res = -1;
+
     if (globals != NULL &&
         PyMapping_GetOptionalItem(globals, &_Py_ID(__lazy_modules__), &lazy_modules) < 0) {
         return -1;
@@ -3047,15 +3050,19 @@ check_lazy_import_comatibility(PyThreadState *tstate, PyObject *globals,
 
     int ilevel = PyLong_AsInt(level);
     if (ilevel == -1 && _PyErr_Occurred(tstate)) {
-        return -1;
+        goto error;
     }
 
-    PyObject *abs_name = _PyImport_GetAbsName(tstate, name, globals, ilevel);
+    abs_name = _PyImport_GetAbsName(tstate, name, globals, ilevel);
     if (abs_name == NULL) {
-        return -1;
+        goto error;
     }
 
-    return PySequence_Contains(lazy_modules, abs_name);
+    res = PySequence_Contains(lazy_modules, abs_name);
+error:
+    Py_XDECREF(abs_name);
+    Py_XDECREF(lazy_modules);
+    return res;
 }
 
 PyObject *
@@ -3089,14 +3096,6 @@ _PyEval_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *glob
         return _PyEval_ImportName(tstate, builtins, globals, locals, name, fromlist, level);
     }
 
-    PyObject *import_func;
-    if (PyMapping_GetOptionalItem(builtins, &_Py_ID(__import__), &import_func) < 0) {
-        goto error;
-    } else if (import_func == NULL) {
-        _PyErr_SetString(tstate, PyExc_ImportError, "__import__ not found");
-        goto error;
-    }
-
     PyObject *lazy_import_func;
     if (PyMapping_GetOptionalItem(builtins, &_Py_ID(__lazy_import__), &lazy_import_func) < 0) {
         goto error;
@@ -3116,16 +3115,15 @@ _PyEval_LazyImportName(PyThreadState *tstate, PyObject *builtins, PyObject *glob
         }
 
         res = _PyImport_LazyImportModuleLevelObject(
-            tstate, name, import_func, globals, locals, fromlist, ilevel
+            tstate, name, builtins, globals, locals, fromlist, ilevel
         );
         goto error;
     }
 
-    PyObject* args[6] = {name, globals, locals, fromlist, level, import_func};
+    PyObject* args[6] = {name, globals, locals, fromlist, level, builtins};
     res = PyObject_Vectorcall(lazy_import_func, args, 6, NULL);
 error:
     Py_XDECREF(lazy_import_func);
-    Py_XDECREF(import_func);
     return res;
 }
 
@@ -3311,8 +3309,10 @@ _PyEval_LazyImportFrom(PyThreadState *tstate, PyObject *v, PyObject *name)
             PyObject *mod_dict = PyModule_GetDict(mod);
             if (mod_dict != NULL) {
                 if (PyDict_GetItemRef(mod_dict, name, &ret) < 0) {
+                    Py_DECREF(mod);
                     return NULL;
                 } else if (ret != NULL) {
+                    Py_DECREF(mod);
                     return ret;
                 }
             }
@@ -3323,20 +3323,20 @@ _PyEval_LazyImportFrom(PyThreadState *tstate, PyObject *v, PyObject *name)
     if (d->lz_attr != NULL) {
         if (PyUnicode_Check(d->lz_attr)) {
             PyObject *from = PyUnicode_FromFormat("%U.%U", d->lz_from, d->lz_attr);
-            ret = _PyLazyImport_New(d->lz_import_func, from, name);
+            ret = _PyLazyImport_New(d->lz_builtins, from, name);
             Py_DECREF(from);
             return ret;
         }
     } else {
         Py_ssize_t dot = PyUnicode_FindChar(d->lz_from, '.', 0, PyUnicode_GET_LENGTH(d->lz_from), 1);
         if (dot >= 0) {
             PyObject *from = PyUnicode_Substring(d->lz_from, 0, dot);
-            ret = _PyLazyImport_New(d->lz_import_func, from, name);
+            ret = _PyLazyImport_New(d->lz_builtins, from, name);
             Py_DECREF(from);
             return ret;
         }
     }
-    ret = _PyLazyImport_New(d->lz_import_func, d->lz_from, name);
+    ret = _PyLazyImport_New(d->lz_builtins, d->lz_from, name);
     return ret;
 }
 
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
diff --git a/Python/import.c b/Python/import.c
diff --git a/Python/specialize.c b/Python/specialize.c

Original file line number	Diff line number	Diff line change
`@@ -58,6 +58,8 @@ extern Py_ssize_t _PyModule_GetFilenameUTF8(`
`58`	`58`	`PyObject* _Py_module_getattro_impl(PyModuleObject m, PyObject name, int suppress);`
`59`	`59`	`PyObject* _Py_module_getattro(PyObject m, PyObject name);`
`60`	`60`
	`61`	`+PyAPI_FUNC(int) _PyModule_ReplaceLazyValue(PyObject dict, PyObject name, PyObject *value);`
	`62`	`+`
`61`	`63`	`#ifdef __cplusplus`
`62`	`64`	`}`
`63`	`65`	`#endif`
Original file line number	Diff line number	Diff line change
`@@ -4719,6 +4719,14 @@ sizeof_lock_held(PyDictObject *mp)`
`4719`	`4719`	`return (Py_ssize_t)res;`
`4720`	`4720`	`}`
`4721`	`4721`
	`4722`	`+void`
	`4723`	`+_PyDict_ClearKeysVersion(PyObject *mp)`
	`4724`	`+{`
	`4725`	`+ ASSERT_DICT_LOCKED(mp);`
	`4726`	`+`
	`4727`	`+ FT_ATOMIC_STORE_UINT32_RELAXED(((PyDictObject *)mp)->ma_keys->dk_version, 0);`
	`4728`	`+}`
	`4729`	`+`
`4722`	`4730`	`Py_ssize_t`
`4723`	`4731`	`_PyDict_SizeOf(PyDictObject *mp)`
`4724`	`4732`	`{`