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

11 files changed, 235 insertions, 14 deletions

diff --git a/Include/cpython/pyerrors.h b/Include/cpython/pyerrors.h
index f33d3ca..1413416 100644
--- a/Include/cpython/pyerrors.h
+++ b/Include/cpython/pyerrors.h
@@ -37,6 +37,7 @@ typedef struct {
     PyObject *msg;
     PyObject *name;
     PyObject *path;
+    PyObject *name_from;
 } PyImportErrorObject;
 
 typedef struct {
@@ -176,4 +177,11 @@ PyAPI_FUNC(void) _Py_NO_RETURN _Py_FatalErrorFormat(
     const char *format,
     ...);
 
+extern PyObject *_PyErr_SetImportErrorWithNameFrom(
+        PyObject *,
+        PyObject *,
+        PyObject *,
+        PyObject *);
+
+
 #define Py_FatalError(message) _Py_FatalErrorFunc(__func__, (message))
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
index 9c716a3..43f4dac 100644
--- a/Include/internal/pycore_global_strings.h
+++ b/Include/internal/pycore_global_strings.h
@@ -478,6 +478,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(n_sequence_fields)
         STRUCT_FOR_ID(n_unnamed_fields)
         STRUCT_FOR_ID(name)
+        STRUCT_FOR_ID(name_from)
         STRUCT_FOR_ID(namespace_separator)
         STRUCT_FOR_ID(namespaces)
         STRUCT_FOR_ID(narg)
diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h
index 55c7c9e..f7823d3 100644
--- a/Include/internal/pycore_runtime_init_generated.h
+++ b/Include/internal/pycore_runtime_init_generated.h
@@ -987,6 +987,7 @@ extern "C" {
                 INIT_ID(n_sequence_fields), \
                 INIT_ID(n_unnamed_fields), \
                 INIT_ID(name), \
+                INIT_ID(name_from), \
                 INIT_ID(namespace_separator), \
                 INIT_ID(namespaces), \
                 INIT_ID(narg), \
@@ -2286,6 +2287,8 @@ _PyUnicode_InitStaticStrings(void) {
     PyUnicode_InternInPlace(&string);
     string = &_Py_ID(name);
     PyUnicode_InternInPlace(&string);
+    string = &_Py_ID(name_from);
+    PyUnicode_InternInPlace(&string);
     string = &_Py_ID(namespace_separator);
     PyUnicode_InternInPlace(&string);
     string = &_Py_ID(namespaces);
@@ -6505,6 +6508,10 @@ _PyStaticObjects_CheckRefcnt(void) {
         _PyObject_Dump((PyObject *)&_Py_ID(name));
         Py_FatalError("immortal object has less refcnt than expected _PyObject_IMMORTAL_REFCNT");
     };
+    if (Py_REFCNT((PyObject *)&_Py_ID(name_from)) < _PyObject_IMMORTAL_REFCNT) {
+        _PyObject_Dump((PyObject *)&_Py_ID(name_from));
+        Py_FatalError("immortal object has less refcnt than expected _PyObject_IMMORTAL_REFCNT");
+    };
     if (Py_REFCNT((PyObject *)&_Py_ID(namespace_separator)) < _PyObject_IMMORTAL_REFCNT) {
         _PyObject_Dump((PyObject *)&_Py_ID(namespace_separator));
         Py_FatalError("immortal object has less refcnt than expected _PyObject_IMMORTAL_REFCNT");
diff --git a/Lib/test/test_call.py b/Lib/test/test_call.py
index 1f3307f..0b37116 100644
--- a/Lib/test/test_call.py
+++ b/Lib/test/test_call.py
@@ -140,9 +140,9 @@ class CFunctionCallsErrorMessages(unittest.TestCase):
                                itertools.product, 0, repeat=1, foo=2)
 
     def test_varargs15_kw(self):
-        msg = r"^ImportError\(\) takes at most 2 keyword arguments \(3 given\)$"
+        msg = r"^ImportError\(\) takes at most 3 keyword arguments \(4 given\)$"
         self.assertRaisesRegex(TypeError, msg,
-                               ImportError, 0, name=1, path=2, foo=3)
+                               ImportError, 0, name=1, path=2, name_from=3, foo=3)
 
     def test_varargs16_kw(self):
         msg = r"^min\(\) takes at most 2 keyword arguments \(3 given\)$"
diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
index 2d17e06..cf52d17 100644
--- a/Lib/test/test_traceback.py
+++ b/Lib/test/test_traceback.py
@@ -6,14 +6,20 @@ import linecache
 import sys
 import types
 import inspect
+import importlib
 import unittest
 import re
+import tempfile
+import random
+import string
 from test import support
+import shutil
 from test.support import (Error, captured_output, cpython_only, ALWAYS_EQ,
                           requires_debug_ranges, has_no_debug_ranges,
                           requires_subprocess)
 from test.support.os_helper import TESTFN, unlink
 from test.support.script_helper import assert_python_ok, assert_python_failure
+from test.support.import_helper import forget
 
 import json
 import textwrap
@@ -2985,6 +2991,122 @@ class SuggestionFormattingTestBase:
         self.assertIn("Did you mean", actual)
         self.assertIn("bluch", actual)
 
+    def make_module(self, code):
+        tmpdir = Path(tempfile.mkdtemp())
+        self.addCleanup(shutil.rmtree, tmpdir)
+
+        sys.path.append(str(tmpdir))
+        self.addCleanup(sys.path.pop)
+
+        mod_name = ''.join(random.choices(string.ascii_letters, k=16))
+        module = tmpdir / (mod_name + ".py")
+        module.write_text(code)
+
+        return mod_name
+
+    def get_import_from_suggestion(self, mod_dict, name):
+        modname = self.make_module(mod_dict)
+
+        def callable():
+            try:
+                exec(f"from {modname} import {name}")
+            except ImportError as e:
+                raise e from None
+            except Exception as e:
+                self.fail(f"Expected ImportError but got {type(e)}")
+        self.addCleanup(forget, modname)
+
+        result_lines = self.get_exception(
+            callable, slice_start=-1, slice_end=None
+        )
+        return result_lines[0]
+
+    def test_import_from_suggestions(self):
+        substitution = textwrap.dedent("""\
+            noise = more_noise = a = bc = None
+            blech = None
+        """)
+
+        elimination = textwrap.dedent("""
+            noise = more_noise = a = bc = None
+            blch = None
+        """)
+
+        addition = textwrap.dedent("""
+            noise = more_noise = a = bc = None
+            bluchin = None
+        """)
+
+        substitutionOverElimination = textwrap.dedent("""
+            blach = None
+            bluc = None
+        """)
+
+        substitutionOverAddition = textwrap.dedent("""
+            blach = None
+            bluchi = None
+        """)
+
+        eliminationOverAddition = textwrap.dedent("""
+            blucha = None
+            bluc = None
+        """)
+
+        caseChangeOverSubstitution = textwrap.dedent("""
+            Luch = None
+            fluch = None
+            BLuch = None
+        """)
+
+        for code, suggestion in [
+            (addition, "'bluchin'?"),
+            (substitution, "'blech'?"),
+            (elimination, "'blch'?"),
+            (addition, "'bluchin'?"),
+            (substitutionOverElimination, "'blach'?"),
+            (substitutionOverAddition, "'blach'?"),
+            (eliminationOverAddition, "'bluc'?"),
+            (caseChangeOverSubstitution, "'BLuch'?"),
+        ]:
+            actual = self.get_import_from_suggestion(code, 'bluch')
+            self.assertIn(suggestion, actual)
+
+    def test_import_from_suggestions_do_not_trigger_for_long_attributes(self):
+        code = "blech = None"
+
+        actual = self.get_suggestion(code, 'somethingverywrong')
+        self.assertNotIn("blech", actual)
+
+    def test_import_from_error_bad_suggestions_do_not_trigger_for_small_names(self):
+        code = "vvv = mom = w = id = pytho = None"
+
+        for name in ("b", "v", "m", "py"):
+            with self.subTest(name=name):
+                actual = self.get_import_from_suggestion(code, name)
+                self.assertNotIn("you mean", actual)
+                self.assertNotIn("vvv", actual)
+                self.assertNotIn("mom", actual)
+                self.assertNotIn("'id'", actual)
+                self.assertNotIn("'w'", actual)
+                self.assertNotIn("'pytho'", actual)
+
+    def test_import_from_suggestions_do_not_trigger_for_big_namespaces(self):
+        # A module with lots of names will not be considered for suggestions.
+        chunks = [f"index_{index} = " for index in range(200)]
+        chunks.append(" None")
+        code = " ".join(chunks)
+        actual = self.get_import_from_suggestion(code, 'bluch')
+        self.assertNotIn("blech", actual)
+
+    def test_import_from_error_with_bad_name(self):
+        def raise_attribute_error_with_bad_name():
+            raise ImportError(name=12, obj=23, name_from=11)
+
+        result_lines = self.get_exception(
+            raise_attribute_error_with_bad_name, slice_start=-1, slice_end=None
+        )
+        self.assertNotIn("?", result_lines[-1])
+
     def test_name_error_suggestions(self):
         def Substitution():
             noise = more_noise = a = bc = None
diff --git a/Lib/traceback.py b/Lib/traceback.py
index bb7856a..6270100 100644
--- a/Lib/traceback.py
+++ b/Lib/traceback.py
@@ -707,9 +707,16 @@ class TracebackException:
             self.offset = exc_value.offset
             self.end_offset = exc_value.end_offset
             self.msg = exc_value.msg
+        elif exc_type and issubclass(exc_type, ImportError) and \
+                getattr(exc_value, "name_from", None) is not None:
+            wrong_name = getattr(exc_value, "name_from", None)
+            suggestion = _compute_suggestion_error(exc_value, exc_traceback, wrong_name)
+            if suggestion:
+                self._str += f". Did you mean: '{suggestion}'?"
         elif exc_type and issubclass(exc_type, (NameError, AttributeError)) and \
                 getattr(exc_value, "name", None) is not None:
-            suggestion = _compute_suggestion_error(exc_value, exc_traceback)
+            wrong_name = getattr(exc_value, "name", None)
+            suggestion = _compute_suggestion_error(exc_value, exc_traceback, wrong_name)
             if suggestion:
                 self._str += f". Did you mean: '{suggestion}'?"
             if issubclass(exc_type, NameError):
@@ -1005,8 +1012,7 @@ def _substitution_cost(ch_a, ch_b):
     return _MOVE_COST
 
 
-def _compute_suggestion_error(exc_value, tb):
-    wrong_name = getattr(exc_value, "name", None)
+def _compute_suggestion_error(exc_value, tb, wrong_name):
     if wrong_name is None or not isinstance(wrong_name, str):
         return None
     if isinstance(exc_value, AttributeError):
@@ -1015,6 +1021,12 @@ def _compute_suggestion_error(exc_value, tb):
             d = dir(obj)
         except Exception:
             return None
+    elif isinstance(exc_value, ImportError):
+        try:
+            mod = __import__(exc_value.name)
+            d = dir(mod)
+        except Exception:
+            return None
     else:
         assert isinstance(exc_value, NameError)
         # find most recent frame
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-10-15-22-25-20.gh-issue-91058.Uo2kW-.rst b/Misc/NEWS.d/next/Core and Builtins/2022-10-15-22-25-20.gh-issue-91058.Uo2kW-.rst
new file mode 100644
index 0000000..042c7ce
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-10-15-22-25-20.gh-issue-91058.Uo2kW-.rst
@@ -0,0 +1,3 @@
+:exc:`ImportError` raised from failed ``from <module> import <name>`` now
+include suggestions for the value of ``<name>`` based on the available names
+in ``<module>``. Patch by Pablo Galindo
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 80e98bb..4b4f31a 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -1464,11 +1464,12 @@ SimpleExtendsException(PyExc_BaseException, KeyboardInterrupt,
 static int
 ImportError_init(PyImportErrorObject *self, PyObject *args, PyObject *kwds)
 {
-    static char *kwlist[] = {"name", "path", 0};
+    static char *kwlist[] = {"name", "path", "name_from", 0};
     PyObject *empty_tuple;
     PyObject *msg = NULL;
     PyObject *name = NULL;
     PyObject *path = NULL;
+    PyObject *name_from = NULL;
 
     if (BaseException_init((PyBaseExceptionObject *)self, args, NULL) == -1)
         return -1;
@@ -1476,8 +1477,8 @@ ImportError_init(PyImportErrorObject *self, PyObject *args, PyObject *kwds)
     empty_tuple = PyTuple_New(0);
     if (!empty_tuple)
         return -1;
-    if (!PyArg_ParseTupleAndKeywords(empty_tuple, kwds, "|$OO:ImportError", kwlist,
-                                     &name, &path)) {
+    if (!PyArg_ParseTupleAndKeywords(empty_tuple, kwds, "|$OOO:ImportError", kwlist,
+                                     &name, &path, &name_from)) {
         Py_DECREF(empty_tuple);
         return -1;
     }
@@ -1489,6 +1490,9 @@ ImportError_init(PyImportErrorObject *self, PyObject *args, PyObject *kwds)
     Py_XINCREF(path);
     Py_XSETREF(self->path, path);
 
+    Py_XINCREF(name_from);
+    Py_XSETREF(self->name_from, name_from);
+
     if (PyTuple_GET_SIZE(args) == 1) {
         msg = PyTuple_GET_ITEM(args, 0);
         Py_INCREF(msg);
@@ -1504,6 +1508,7 @@ ImportError_clear(PyImportErrorObject *self)
     Py_CLEAR(self->msg);
     Py_CLEAR(self->name);
     Py_CLEAR(self->path);
+    Py_CLEAR(self->name_from);
     return BaseException_clear((PyBaseExceptionObject *)self);
 }
 
@@ -1521,6 +1526,7 @@ ImportError_traverse(PyImportErrorObject *self, visitproc visit, void *arg)
     Py_VISIT(self->msg);
     Py_VISIT(self->name);
     Py_VISIT(self->path);
+    Py_VISIT(self->name_from);
     return BaseException_traverse((PyBaseExceptionObject *)self, visit, arg);
 }
 
@@ -1540,7 +1546,7 @@ static PyObject *
 ImportError_getstate(PyImportErrorObject *self)
 {
     PyObject *dict = ((PyBaseExceptionObject *)self)->dict;
-    if (self->name || self->path) {
+    if (self->name || self->path || self->name_from) {
         dict = dict ? PyDict_Copy(dict) : PyDict_New();
         if (dict == NULL)
             return NULL;
@@ -1552,6 +1558,10 @@ ImportError_getstate(PyImportErrorObject *self)
             Py_DECREF(dict);
             return NULL;
         }
+        if (self->name_from && PyDict_SetItem(dict, &_Py_ID(name_from), self->name_from) < 0) {
+            Py_DECREF(dict);
+            return NULL;
+        }
         return dict;
     }
     else if (dict) {
@@ -1588,6 +1598,8 @@ static PyMemberDef ImportError_members[] = {
         PyDoc_STR("module name")},
     {"path", T_OBJECT, offsetof(PyImportErrorObject, path), 0,
         PyDoc_STR("module path")},
+    {"name_from", T_OBJECT, offsetof(PyImportErrorObject, name_from), 0,
+        PyDoc_STR("name imported from module")},
     {NULL}  /* Sentinel */
 };
 
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;
 }