Chain lazy IMPORT_FROM placeholders instead of flagging dotted imports

Each deferred IMPORT_FROM off a placeholder without a fromlist now keeps
the previous placeholder in lz_from and the attribute name in lz_attr.
Reification walks back to the placeholder IMPORT_NAME left, runs that
import, and replays the lookups in order with _PyEval_ImportFrom, which
is what the eager bytecode does.  This drops the lz_dotted_as flag and
also follows a custom __lazy_import__ that returns a placeholder for a
different module name.

Author: brittanyrey

Include/internal/pycore_lazyimportobject.h

@@ -22,14 +22,12 @@ typedef struct {
     // 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.
-    int lz_dotted_as;          // `import a.b as c`, which binds `a.b` not `a`.
 } PyLazyImportObject;
 
 
 PyAPI_FUNC(PyObject *) _PyLazyImport_GetName(PyObject *lazy_import);
 PyAPI_FUNC(PyObject *) _PyLazyImport_New(
-    struct _PyInterpreterFrame *frame, PyObject *import_func, PyObject *from,
-    PyObject *attr, int dotted_as);
+    struct _PyInterpreterFrame *frame, PyObject *import_func, PyObject *from, PyObject *attr);
 
 #ifdef __cplusplus
 }

Lib/test/test_lazy_import/__init__.py

@@ -2234,6 +2234,36 @@ def test_lazy_import_as_wins_over_variable(self):
         """)
         assert_python_ok("-c", code)
 
+    def test_dotted_as_of_loaded_module(self):
+        """A dotted lazy import as binds the module, not a same-named attribute."""
+        # importlib.metadata is already loaded and has a `metadata` attribute.
+        code = textwrap.dedent("""
+            import importlib.metadata
+            import importlib.metadata as eager
+
+            lazy import importlib.metadata as lazily
+
+            assert lazily is eager, lazily
+        """)
+        assert_python_ok("-c", code)
+
+    def test_dotted_as_replays_lookups_on_custom_placeholder(self):
+        """A dotted lazy import as looks up its names on what the hook returned."""
+        code = textwrap.dedent("""
+            import builtins
+            import xml.dom
+
+            # In a list, so the hook reading it does not resolve it.
+            placeholder = [__lazy_import__("xml")]
+            default = builtins.__lazy_import__
+            builtins.__lazy_import__ = lambda *args: placeholder[0]
+            lazy import fake.dom as dom
+            builtins.__lazy_import__ = default
+
+            assert dom is xml.dom, dom
+        """)
+        assert_python_ok("-c", code)
+
 
 class DeletedModuleReimportTests(unittest.TestCase):
     """Tests for reimporting after module deletion from sys.modules."""

Objects/lazyimportobject.c

@@ -11,12 +11,11 @@
 #define PyLazyImportObject_CAST(op) ((PyLazyImportObject *)(op))
 
 PyObject *
-_PyLazyImport_New(_PyInterpreterFrame *frame, PyObject *builtins,
-                  PyObject *name, PyObject *fromlist, int dotted_as)
+_PyLazyImport_New(_PyInterpreterFrame *frame, PyObject *builtins, PyObject *name, PyObject *fromlist)
 {
     PyLazyImportObject *m;
-    if (!name || !PyUnicode_Check(name)) {
-        PyErr_SetString(PyExc_TypeError, "expected str for name");
+    if (!name || !(PyUnicode_Check(name) || PyLazyImport_CheckExact(name))) {
+        PyErr_SetString(PyExc_TypeError, "expected str or lazy_import for name");
         return NULL;
     }
     if (fromlist == Py_None || fromlist == NULL) {
@@ -34,7 +33,6 @@ _PyLazyImport_New(_PyInterpreterFrame *frame, PyObject *builtins,
     m->lz_builtins = Py_XNewRef(builtins);
     m->lz_from = Py_NewRef(name);
     m->lz_attr = Py_XNewRef(fromlist);
-    m->lz_dotted_as = dotted_as;
 
     // Capture frame information for the original import location.
     m->lz_code = NULL;
@@ -106,9 +104,43 @@ lazy_import_getattro(PyObject *op, PyObject *name)
     return value;
 }
 
+static PyObject *lazy_import_name(PyLazyImportObject *m);
+
+// The dotted name of the object that resolving the placeholder returns.
+static PyObject *
+lazy_import_path(PyLazyImportObject *m)
+{
+    if (PyLazyImport_CheckExact(m->lz_from)) {
+        PyObject *base = lazy_import_path((PyLazyImportObject *)m->lz_from);
+        if (base == NULL) {
+            return NULL;
+        }
+        PyObject *res = PyUnicode_FromFormat("%U.%U", base, m->lz_attr);
+        Py_DECREF(base);
+        return res;
+    }
+    if (m->lz_attr != NULL) {
+        return lazy_import_name(m);
+    }
+    // __import__("a.b") returns the top-level package `a`.
+    Py_ssize_t dot = PyUnicode_FindChar(
+        m->lz_from, '.', 0, PyUnicode_GET_LENGTH(m->lz_from), 1
+    );
+    if (dot == -2) {
+        return NULL;
+    }
+    if (dot < 0) {
+        return Py_NewRef(m->lz_from);
+    }
+    return PyUnicode_Substring(m->lz_from, 0, dot);
+}
+
 static PyObject *
 lazy_import_name(PyLazyImportObject *m)
 {
+    if (PyLazyImport_CheckExact(m->lz_from)) {
+        return lazy_import_path(m);
+    }
     if (m->lz_attr != NULL) {
         if (PyUnicode_Check(m->lz_attr)) {
             return PyUnicode_FromFormat("%U.%U", m->lz_from, m->lz_attr);

Python/ceval.c

@@ -3334,12 +3334,10 @@ _PyEval_LazyImportFrom(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObje
     PyObject *ret;
     PyLazyImportObject *d = (PyLazyImportObject *)v;
 
-    if (d->lz_attr == NULL) {
-        // `import a.b as x` binds the module `a.b`, not an attribute of `a`.
-        if (d->lz_dotted_as) {
-            return Py_NewRef(v);
-        }
-        return _PyLazyImport_New(frame, d->lz_builtins, d->lz_from, NULL, 1);
+    if (d->lz_attr == NULL || PyUnicode_Check(d->lz_attr)) {
+        // `import a.b.c as d`: record the lookup, to replay once the
+        // import has run.
+        return _PyLazyImport_New(frame, d->lz_builtins, v, name);
     }
 
     PyObject *mod = PyImport_GetModule(d->lz_from);
@@ -3362,8 +3360,7 @@ _PyEval_LazyImportFrom(PyThreadState *tstate, _PyInterpreterFrame *frame, PyObje
         Py_DECREF(mod);
     }
 
-    assert(!PyUnicode_Check(d->lz_attr));  // a fromlist
-    return _PyLazyImport_New(frame, d->lz_builtins, d->lz_from, name, 0);
+    return _PyLazyImport_New(frame, d->lz_builtins, d->lz_from, name);
 }
 
 #define CANNOT_CATCH_MSG "catching classes that do not inherit from "\

Python/import.c

@@ -3898,23 +3898,28 @@ _PyImport_ResolveName(PyThreadState *tstate, PyObject *name,
   return resolve_name(tstate, name, globals, level);
 }
 
-// Take `a.b.c` off the `a` that importing "a.b.c" returns, as IMPORT_FROM does.
+// Look up, in order, the attributes recorded from the root placeholder to lz
+// on the module the root's import returned.
 static PyObject *
-import_from_dotted_name(PyThreadState *tstate, PyObject *mod, PyObject *name)
-{
-    PyObject *parts = PyUnicode_Split(name, _Py_LATIN1_CHR('.'), -1);
-    if (parts == NULL) {
-        return NULL;
-    }
-    PyObject *obj = Py_NewRef(mod);
-    for (Py_ssize_t i = 1; i < PyList_GET_SIZE(parts); i++) {
-        Py_SETREF(obj, _PyEval_ImportFrom(tstate, obj,
-                                          PyList_GET_ITEM(parts, i)));
-        if (obj == NULL) {
-            break;
+lazy_import_replay_from(PyThreadState *tstate, PyObject *mod,
+                        PyLazyImportObject *lz)
+{
+    PyObject *from;
+    if (PyLazyImport_CheckExact(lz->lz_from)) {
+        from = lazy_import_replay_from(
+            tstate, mod, (PyLazyImportObject *)lz->lz_from);
+        if (from == NULL) {
+            return NULL;
         }
     }
-    Py_DECREF(parts);
+    else if (lz->lz_attr != NULL && PyUnicode_Check(lz->lz_attr)) {
+        from = Py_NewRef(mod);
+    }
+    else {
+        return Py_NewRef(mod);
+    }
+    PyObject *obj = _PyEval_ImportFrom(tstate, from, lz->lz_attr);
+    Py_DECREF(from);
     return obj;
 }
 
@@ -3930,6 +3935,12 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     PyLazyImportObject *lz = (PyLazyImportObject *)lazy_import;
     PyInterpreterState *interp = tstate->interp;
 
+    // Walk back to the placeholder IMPORT_NAME left.
+    PyLazyImportObject *root = lz;
+    while (PyLazyImport_CheckExact(root->lz_from)) {
+        root = (PyLazyImportObject *)root->lz_from;
+    }
+
     // Acquire the global import lock to serialize reification
     _PyImport_AcquireLock(interp);
 
@@ -3966,7 +3977,7 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
             return NULL;
         }
         PyErr_SetImportErrorSubclass(PyExc_ImportCycleError, errmsg,
-                                     lz->lz_from, NULL);
+                                     root->lz_from, NULL);
         Py_DECREF(errmsg);
         Py_DECREF(name);
         _PyImport_ReleaseLock(interp);
@@ -3976,24 +3987,24 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
         goto error;
     }
 
-    if (lz->lz_attr != NULL) {
-        if (PyUnicode_Check(lz->lz_attr)) {
+    if (root->lz_attr != NULL) {
+        if (PyUnicode_Check(root->lz_attr)) {
             fromlist = PyTuple_New(1);
             if (fromlist == NULL) {
                 goto error;
             }
-            Py_INCREF(lz->lz_attr);
-            PyTuple_SET_ITEM(fromlist, 0, lz->lz_attr);
+            Py_INCREF(root->lz_attr);
+            PyTuple_SET_ITEM(fromlist, 0, root->lz_attr);
         }
         else {
-            Py_INCREF(lz->lz_attr);
-            fromlist = lz->lz_attr;
+            Py_INCREF(root->lz_attr);
+            fromlist = root->lz_attr;
         }
     }
 
     PyObject *globals = PyEval_GetGlobals();
 
-    if (PyMapping_GetOptionalItem(lz->lz_builtins, &_Py_ID(__import__),
+    if (PyMapping_GetOptionalItem(root->lz_builtins, &_Py_ID(__import__),
                                   &import_func) < 0) {
         goto error;
     }
@@ -4003,27 +4014,17 @@ _PyImport_LoadLazyImportTstate(PyThreadState *tstate, PyObject *lazy_import)
     }
     obj = _PyEval_ImportNameWithImport(
         tstate, import_func, globals, globals,
-        lz->lz_from, fromlist, _PyLong_GetZero()
+        root->lz_from, fromlist, _PyLong_GetZero()
     );
     if (obj == NULL) {
         goto error;
     }
 
-    if (lz->lz_dotted_as) {
-        PyObject *top = obj;
-        obj = import_from_dotted_name(tstate, top, lz->lz_from);
-        Py_DECREF(top);
-        if (obj == NULL) {
-            goto error;
-        }
-    }
-    else if (lz->lz_attr != NULL && PyUnicode_Check(lz->lz_attr)) {
-        PyObject *from = obj;
-        obj = _PyEval_ImportFrom(tstate, from, lz->lz_attr);
-        Py_DECREF(from);
-        if (obj == NULL) {
-            goto error;
-        }
+    PyObject *from = obj;
+    obj = lazy_import_replay_from(tstate, from, lz);
+    Py_DECREF(from);
+    if (obj == NULL) {
+        goto error;
     }
 
     assert(!PyLazyImport_CheckExact(obj));
@@ -4628,7 +4629,7 @@ _PyImport_LazyImportModuleLevelObject(PyThreadState *tstate,
     else {
         Py_XINCREF(fromlist);
     }
-    PyObject *res = _PyLazyImport_New(frame, builtins, abs_name, fromlist, 0);
+    PyObject *res = _PyLazyImport_New(frame, builtins, abs_name, fromlist);
     if (res == NULL) {
         Py_XDECREF(fromlist);
         Py_DECREF(abs_name);