gh-91058: Add error suggestions to 'import from' import errors (#98305)

3 files changed, 61 insertions, 5 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index fb8dd48..35ce767 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -6900,7 +6900,7 @@ import_from(PyThreadState *tstate, PyObject *v, PyObject *name)
             name, pkgname_or_unknown
         );
         /* NULL checks for errmsg and pkgname done by PyErr_SetImportError. */
-        PyErr_SetImportError(errmsg, pkgname, NULL);
+        _PyErr_SetImportErrorWithNameFrom(errmsg, pkgname, NULL, name);
     }
     else {
         PyObject *spec = PyObject_GetAttr(v, &_Py_ID(__spec__));
@@ -6913,7 +6913,7 @@ import_from(PyThreadState *tstate, PyObject *v, PyObject *name)
 
         errmsg = PyUnicode_FromFormat(fmt, name, pkgname_or_unknown, pkgpath);
         /* NULL checks for errmsg and pkgname done by PyErr_SetImportError. */
-        PyErr_SetImportError(errmsg, pkgname, pkgpath);
+        _PyErr_SetImportErrorWithNameFrom(errmsg, pkgname, pkgpath, name);
     }
 
     Py_XDECREF(errmsg);
diff --git a/Python/errors.c b/Python/errors.c
index 2aa748c..fc3e468 100644
--- a/Python/errors.c
+++ b/Python/errors.c
@@ -975,9 +975,10 @@ PyObject *PyErr_SetFromWindowsErrWithFilename(
 
 #endif /* MS_WINDOWS */
 
-PyObject *
-PyErr_SetImportErrorSubclass(PyObject *exception, PyObject *msg,
-    PyObject *name, PyObject *path)
+static PyObject *
+_PyErr_SetImportErrorSubclassWithNameFrom(
+    PyObject *exception, PyObject *msg,
+    PyObject *name, PyObject *path, PyObject* from_name)
 {
     PyThreadState *tstate = _PyThreadState_GET();
     int issubclass;
@@ -1005,6 +1006,10 @@ PyErr_SetImportErrorSubclass(PyObject *exception, PyObject *msg,
     if (path == NULL) {
         path = Py_None;
     }
+    if (from_name == NULL) {
+        from_name = Py_None;
+    }
+
 
     kwargs = PyDict_New();
     if (kwargs == NULL) {
@@ -1016,6 +1021,9 @@ PyErr_SetImportErrorSubclass(PyObject *exception, PyObject *msg,
     if (PyDict_SetItemString(kwargs, "path", path) < 0) {
         goto done;
     }
+    if (PyDict_SetItemString(kwargs, "name_from", from_name) < 0) {
+        goto done;
+    }
 
     error = PyObject_VectorcallDict(exception, &msg, 1, kwargs);
     if (error != NULL) {
@@ -1028,6 +1036,20 @@ done:
     return NULL;
 }
 
+
+PyObject *
+PyErr_SetImportErrorSubclass(PyObject *exception, PyObject *msg,
+    PyObject *name, PyObject *path)
+{
+    return _PyErr_SetImportErrorSubclassWithNameFrom(exception, msg, name, path, NULL);
+}
+
+PyObject *
+_PyErr_SetImportErrorWithNameFrom(PyObject *msg, PyObject *name, PyObject *path, PyObject* from_name)
+{
+    return _PyErr_SetImportErrorSubclassWithNameFrom(PyExc_ImportError, msg, name, path, from_name);
+}
+
 PyObject *
 PyErr_SetImportError(PyObject *msg, PyObject *name, PyObject *path)
 {
diff --git a/Python/suggestions.c b/Python/suggestions.c
index 89b86f7..82376b6 100644
--- a/Python/suggestions.c
+++ b/Python/suggestions.c
@@ -312,6 +312,38 @@ offer_suggestions_for_name_error(PyNameErrorObject *exc)
     return result;
 }
 
+static PyObject *
+offer_suggestions_for_import_error(PyImportErrorObject *exc)
+{
+    PyObject *mod_name = exc->name; // borrowed reference
+    PyObject *name = exc->name_from; // borrowed reference
+    if (name == NULL || mod_name == NULL || name == Py_None ||
+        !PyUnicode_CheckExact(name) || !PyUnicode_CheckExact(mod_name)) {
+        return NULL;
+    }
+
+    PyObject* mod = PyImport_GetModule(mod_name);
+    if (mod == NULL) {
+        return NULL;
+    }
+
+    PyObject *dir = PyObject_Dir(mod);
+    Py_DECREF(mod);
+    if (dir == NULL) {
+        return NULL;
+    }
+
+    PyObject *suggestion = calculate_suggestions(dir, name);
+    Py_DECREF(dir);
+    if (!suggestion) {
+        return NULL;
+    }
+
+    PyObject* result = PyUnicode_FromFormat(". Did you mean: %R?", suggestion);
+    Py_DECREF(suggestion);
+    return result;
+}
+
 // Offer suggestions for a given exception. Returns a python string object containing the
 // suggestions. This function returns NULL if no suggestion was found or if an exception happened,
 // users must call PyErr_Occurred() to disambiguate.
@@ -324,6 +356,8 @@ _Py_Offer_Suggestions(PyObject *exception)
         result = offer_suggestions_for_attribute_error((PyAttributeErrorObject *) exception);
     } else if (Py_IS_TYPE(exception, (PyTypeObject*)PyExc_NameError)) {
         result = offer_suggestions_for_name_error((PyNameErrorObject *) exception);
+    } else if (Py_IS_TYPE(exception, (PyTypeObject*)PyExc_ImportError)) {
+        result = offer_suggestions_for_import_error((PyImportErrorObject *) exception);
     }
     return result;
 }