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

pablogsal · web-flow · commit 7cfbb49fcd4c · 2022-10-25T23:56:59.000+01:00
diff --git 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
@@ -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
diff --git a/Lib/test/test_call.py b/Lib/test/test_call.py
@@ -140,9 +140,9 @@ def test_varargs14_kw(self):
                                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
@@ -6,14 +6,20 @@
 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 @@ def __getattribute__(self, attr):
         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
@@ -707,9 +707,16 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
             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
@@ -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
@@ -1464,20 +1464,21 @@ 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;
 
     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
@@ -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
diff --git a/Python/suggestions.c b/Python/suggestions.c

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	+:exc:`ImportError` raised from failed ``from <module> import <name>`` now
	`2`	+include suggestions for the value of ``<name>`` based on the available names
	`3`	+in ``<module>``. Patch by Pablo Galindo