gh-142349: Simplify lazy import resolution (#158282)

Co-authored-by: Petr Viktorin <encukou@gmail.com>
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>

Author: 3 people

Include/internal/pycore_dict.h

@@ -23,6 +23,11 @@ PyAPI_FUNC(int) _PyDict_DelItemIf(PyObject *mp, PyObject *key,
                                   int (*predicate)(PyObject *value, void *arg),
                                   void *arg);
 
+// Atomically replace an existing value if it is expected (by identity).
+// Return 1 if replaced, 0 if absent or different, or -1 on error.
+extern int _PyDict_ReplaceItemIf(PyObject *dict, PyObject *key,
+                               PyObject *expected, PyObject *replacement);
+
 // "KnownHash" variants
 // Export for '_asyncio' shared extension
 PyAPI_FUNC(int) _PyDict_SetItem_KnownHash(PyObject *mp, PyObject *key,
@@ -147,8 +152,8 @@ PyAPI_FUNC(Py_ssize_t) _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyOb
  */
 PyAPI_FUNC(Py_ssize_t) _PyDictKeys_StringLookupAndVersion(PyDictKeysObject* dictkeys, PyObject *key, uint32_t *version);
 PyAPI_FUNC(Py_ssize_t) _PyDictKeys_StringLookupSplit(PyDictKeysObject* dictkeys, PyObject *key);
-PyAPI_FUNC(PyObject *)_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObject *);
-PyAPI_FUNC(void) _PyDict_LoadGlobalStackRef(PyDictObject *, PyDictObject *, PyObject *, _PyStackRef *);
+// Return the borrowed source dictionary, or NULL if absent or on error.
+PyAPI_FUNC(PyObject *) _PyDict_LoadGlobalStackRef(PyDictObject *, PyDictObject *, PyObject *, _PyStackRef *);
 
 // Loads the __builtins__ object from the globals dict. Returns a new reference.
 extern PyObject *_PyDict_LoadBuiltinsFromGlobals(PyObject *globals);

Include/internal/pycore_import.h

@@ -32,20 +32,16 @@ extern int _PyImport_FixupBuiltin(
     PyObject *modules
     );
 
-extern PyObject * _PyImport_ResolveName(
-    PyThreadState *tstate, PyObject *name, PyObject *globals, int level);
 extern PyObject * _PyImport_GetAbsName(
     PyThreadState *tstate, PyObject *name, PyObject *globals, int level);
 // Symbol is exported for the JIT on Windows builds.
 PyAPI_FUNC(PyObject *) _PyImport_LoadLazyImportTstate(
     PyThreadState *tstate, PyObject *lazy_import);
-typedef enum {
-    _Py_LAZY_SUBMODULE_ERROR = -1,
-    _Py_LAZY_SUBMODULE_NOT_FOUND = 0,
-    _Py_LAZY_SUBMODULE_LOADED = 1,
-} _PyLazySubmoduleImportResult;
-extern _PyLazySubmoduleImportResult _PyImport_TryLoadLazySubmodule(
-    PyObject *mod_name, PyObject *attr_name, PyObject **result);
+// Bind a pending child on module before removing its registration. Return a
+// new reference, NULL without an exception if absent, or NULL on error.
+// Set recheck_dict if a concurrent lookup may have already bound the child.
+extern PyObject * _PyImport_TryLoadLazySubmodule(
+    PyObject *module, PyObject *attr_name, int *recheck_dict);
 extern PyObject * _PyImport_LazyImportModuleLevelObject(
     PyThreadState *tstate, PyObject *name, PyObject *builtins,
     PyObject *globals, PyObject *locals, PyObject *fromlist, int level);

Include/internal/pycore_interp_structs.h

@@ -352,15 +352,12 @@ struct _import_state {
     PyObject *lazy_import_func;
     int lazy_imports_mode;
     PyObject *lazy_imports_filter;
-    PyObject *lazy_importing_modules;
     // The set stored in sys.lazy_modules if values that have been
     // lazily imported. This value is only for debugging/introspection
     // purposes and is not used by the runtime.
     PyObject *lazy_modules;
     // A dict mapping package names to a set of submodule names that
-    // have been imported lazily from packages which have been imported
-    // lazily. When the package is reified we need to add a
-    // LazyImportObject which refers to the submodule on the module.
+    // may need to be imported when an attribute is missing on the package.
     PyObject *lazy_pending_submodules;
 #ifdef Py_GIL_DISABLED
     PyMutex lazy_mutex;

Include/internal/pycore_lazyimportobject.h

@@ -14,20 +14,18 @@ extern "C" {
 PyAPI_DATA(PyTypeObject) PyLazyImport_Type;
 #define PyLazyImport_CheckExact(op) Py_IS_TYPE((op), &PyLazyImport_Type)
 
-typedef struct {
-    PyObject_HEAD
-    PyObject *lz_builtins;
-    PyObject *lz_from;
-    PyObject *lz_attr;
-    // Frame information for the original import location.
-    PyCodeObject *lz_code;     // Code object where the lazy import was created.
-    int lz_instr_offset;       // Instruction offset where the lazy import was created.
-} PyLazyImportObject;
-
-
-PyAPI_FUNC(PyObject *) _PyLazyImport_GetName(PyObject *lazy_import);
 PyAPI_FUNC(PyObject *) _PyLazyImport_New(
-    struct _PyInterpreterFrame *frame, PyObject *import_func, PyObject *from, PyObject *attr);
+    struct _PyInterpreterFrame *frame, PyObject *builtins,
+    PyObject *name, PyObject *fromlist);
+
+extern int _PyLazyImport_IsResolving(PyThreadState *tstate, PyObject *op);
+
+// Resolve a placeholder and replace its binding if it is unchanged or holds
+// the child module published by the normal importer during resolution.
+// namespace is the source captured during lookup, before resolution runs.
+PyAPI_FUNC(PyObject *) _PyLazyImport_Reify(
+    PyThreadState *tstate, PyObject *placeholder,
+    PyObject *name, PyObject *ns);
 
 #ifdef __cplusplus
 }

Include/internal/pycore_opcode_metadata.h

@@ -51,6 +51,9 @@ typedef struct _PyThreadStateImpl {
     PyObject *asyncio_running_loop; // Strong reference
     PyObject *asyncio_running_task; // Strong reference
 
+    // Reused for the thread's lifetime; owns placeholders while they resolve.
+    PyObject *lazy_imports;
+
     // Distinguishes between yield and return from PyEval_EvalFrame().
     // See gen_send_ex2() in Objects/genobject.c
     enum {

Include/internal/pycore_tstate.h

@@ -13,7 +13,7 @@
 import contextlib
 
 from test import support
-from test.support.script_helper import assert_python_ok
+from test.support.script_helper import assert_python_ok, assert_python_failure
 
 try:
     import _testcapi
@@ -617,6 +617,14 @@ def test_dunder_lazy_import_invalid_arguments(self):
         with self.assertRaises(TypeError):
             __lazy_import__("sys", globals=1)
 
+        code = textwrap.dedent("""
+            __lazy_import__("sys", fromlist=(1, 2, 3))
+        """)
+        result = assert_python_failure("-c", code, NO_COLOR='y')
+        self.assertIn(
+            b"TypeError: Item in ``from list'' must be str, not int",
+            result.err)
+
     def test_dunder_lazy_import_builtins(self):
         """__lazy_import__ should use module's __builtins__ for __import__."""
         from test.test_lazy_import.data import dunder_lazy_import_builtins
@@ -766,6 +774,28 @@ def test_missing_lazy_from_import_shows_chained_traceback(self):
         """)
         assert_python_ok("-c", code)
 
+    @support.subTests('name', (
+        'test.test_lazy_import.data.broken_module_chained_cause',
+        'test.test_lazy_import.data.broken_module_chained_context',
+        'test.test_lazy_import.data.broken_module_chained_suppressed',
+    ))
+    def test_chained_exception_import_shows_notes(self, name):
+        """Accessing missing attribute from lazy from-import should chain errors."""
+        code = textwrap.dedent(f"""
+            lazy import {name}
+
+            try:
+                _ = test
+            except ValueError as e:
+                assert any(
+                    note.startswith("lazy import of '{name}' declared in ")
+                    for note in e.__notes__
+                ), e.__notes__
+            else:
+                raise AssertionError("ImportError was not raised")
+        """)
+        assert_python_ok("-c", code)
+
     def test_reification_retries_on_failure(self):
         """Failed reification should allow retry on subsequent access.
 

Include/internal/pycore_uop_metadata.h

@@ -0,0 +1,3 @@
+# Module that raises an exception with explicit cause during import
+cause = ValueError("Cause of failure")
+raise ValueError("This module always fails to import") from cause

Lib/test/test_lazy_import/__init__.py

@@ -0,0 +1,5 @@
+# Module that raises an exception with context during import
+try:
+    raise ValueError("Cause of failure")
+except:
+    raise ValueError("This module always fails to import")