bpo-38530: Offer suggestions on NameError (GH-25397)

When printing NameError raised by the interpreter, PyErr_Display
will offer suggestions of simmilar variable names in the function that the exception
was raised from:

    >>> schwarzschild_black_hole = None
    >>> schwarschild_black_hole
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    NameError: name 'schwarschild_black_hole' is not defined. Did you mean: schwarzschild_black_hole?

8 files changed, 287 insertions, 10 deletions

diff --git a/Doc/library/exceptions.rst b/Doc/library/exceptions.rst
index 8fdd6ebecf..f4f5c478f2 100644
--- a/Doc/library/exceptions.rst
+++ b/Doc/library/exceptions.rst
@@ -242,6 +242,13 @@ The following exceptions are the exceptions that are usually raised.
    unqualified names.  The associated value is an error message that includes the
    name that could not be found.
 
+   The :attr:`name` attribute can be set using a keyword-only argument to the
+   constructor. When set it represent the name of the variable that was attempted
+   to be accessed.
+
+   .. versionchanged:: 3.10
+      Added the :attr:`name` attribute.
+
 
 .. exception:: NotImplementedError
 
diff --git a/Doc/whatsnew/3.10.rst b/Doc/whatsnew/3.10.rst
index f149d7453b..69697e15d4 100644
--- a/Doc/whatsnew/3.10.rst
+++ b/Doc/whatsnew/3.10.rst
@@ -187,6 +187,23 @@ raised from:
 
 (Contributed by Pablo Galindo in :issue:`38530`.)
 
+NameErrors
+~~~~~~~~~~
+
+When printing :exc:`NameError` raised by the interpreter, :c:func:`PyErr_Display`
+will offer suggestions of simmilar variable names in the function that the exception
+was raised from:
+
+.. code-block:: python
+
+    >>> schwarzschild_black_hole = None
+    >>> schwarschild_black_hole
+    Traceback (most recent call last):
+      File "<stdin>", line 1, in <module>
+    NameError: name 'schwarschild_black_hole' is not defined. Did you mean: schwarzschild_black_hole?
+
+(Contributed by Pablo Galindo in :issue:`38530`.)
+
 PEP 626: Precise line numbers for debugging and other tools
 -----------------------------------------------------------
 
diff --git a/Include/cpython/pyerrors.h b/Include/cpython/pyerrors.h
index a15082e693..9d88e6631f 100644
--- a/Include/cpython/pyerrors.h
+++ b/Include/cpython/pyerrors.h
@@ -64,6 +64,11 @@ typedef struct {
 
 typedef struct {
     PyException_HEAD
+    PyObject *name;
+} PyNameErrorObject;
+
+typedef struct {
+    PyException_HEAD
     PyObject *obj;
     PyObject *name;
 } PyAttributeErrorObject;
diff --git a/Lib/test/test_exceptions.py b/Lib/test/test_exceptions.py
index e1a5ec76d7..4f3c9ab456 100644
--- a/Lib/test/test_exceptions.py
+++ b/Lib/test/test_exceptions.py
@@ -1413,6 +1413,129 @@ class ExceptionTests(unittest.TestCase):
 
             gc_collect()
 
+global_for_suggestions = None
+
+class NameErrorTests(unittest.TestCase):
+    def test_name_error_has_name(self):
+        try:
+            bluch
+        except NameError as exc:
+            self.assertEqual("bluch", exc.name)
+
+    def test_name_error_suggestions(self):
+        def Substitution():
+            noise = more_noise = a = bc = None
+            blech = None
+            print(bluch)
+
+        def Elimination():
+            noise = more_noise = a = bc = None
+            blch = None
+            print(bluch)
+
+        def Addition():
+            noise = more_noise = a = bc = None
+            bluchin = None
+            print(bluch)
+
+        def SubstitutionOverElimination():
+            blach = None
+            bluc = None
+            print(bluch)
+
+        def SubstitutionOverAddition():
+            blach = None
+            bluchi = None
+            print(bluch)
+
+        def EliminationOverAddition():
+            blucha = None
+            bluc = None
+            print(bluch)
+
+        for func, suggestion in [(Substitution, "blech?"),
+                                (Elimination, "blch?"),
+                                (Addition, "bluchin?"),
+                                (EliminationOverAddition, "blucha?"),
+                                (SubstitutionOverElimination, "blach?"),
+                                (SubstitutionOverAddition, "blach?")]:
+            err = None
+            try:
+                func()
+            except NameError as exc:
+                with support.captured_stderr() as err:
+                    sys.__excepthook__(*sys.exc_info())
+            self.assertIn(suggestion, err.getvalue())
+
+    def test_name_error_suggestions_from_globals(self):
+        def func():
+            print(global_for_suggestio)
+        try:
+            func()
+        except NameError as exc:
+            with support.captured_stderr() as err:
+                sys.__excepthook__(*sys.exc_info())
+        self.assertIn("global_for_suggestions?", err.getvalue())
+
+    def test_name_error_suggestions_do_not_trigger_for_long_names(self):
+        def f():
+            somethingverywronghehehehehehe = None
+            print(somethingverywronghe)
+
+        try:
+            f()
+        except NameError as exc:
+            with support.captured_stderr() as err:
+                sys.__excepthook__(*sys.exc_info())
+
+        self.assertNotIn("somethingverywronghehe", err.getvalue())
+
+    def test_name_error_suggestions_do_not_trigger_for_big_dicts(self):
+        def f():
+            # Mutating locals() is unreliable, so we need to do it by hand
+            a1 = a2 = a3 = a4 = a5 = a6 = a7 = a8 = a9 = a10 = a11 = a12 = a13 = \
+            a14 = a15 = a16 = a17 = a18 = a19 = a20 = a21 = a22 = a23 = a24 = a25 = \
+            a26 = a27 = a28 = a29 = a30 = a31 = a32 = a33 = a34 = a35 = a36 = a37 = \
+            a38 = a39 = a40 = a41 = a42 = a43 = a44 = a45 = a46 = a47 = a48 = a49 = \
+            a50 = a51 = a52 = a53 = a54 = a55 = a56 = a57 = a58 = a59 = a60 = a61 = \
+            a62 = a63 = a64 = a65 = a66 = a67 = a68 = a69 = a70 = a71 = a72 = a73 = \
+            a74 = a75 = a76 = a77 = a78 = a79 = a80 = a81 = a82 = a83 = a84 = a85 = \
+            a86 = a87 = a88 = a89 = a90 = a91 = a92 = a93 = a94 = a95 = a96 = a97 = \
+            a98 = a99 = a100 = a101 = a102 = a103 = None
+            print(a0)
+
+        try:
+            f()
+        except NameError as exc:
+            with support.captured_stderr() as err:
+                sys.__excepthook__(*sys.exc_info())
+
+        self.assertNotIn("a10", err.getvalue())
+
+    def test_name_error_with_custom_exceptions(self):
+        def f():
+            blech = None
+            raise NameError()
+
+        try:
+            f()
+        except NameError as exc:
+            with support.captured_stderr() as err:
+                sys.__excepthook__(*sys.exc_info())
+
+        self.assertNotIn("blech", err.getvalue())
+
+        def f():
+            blech = None
+            raise NameError
+
+        try:
+            f()
+        except NameError as exc:
+            with support.captured_stderr() as err:
+                sys.__excepthook__(*sys.exc_info())
+
+        self.assertNotIn("blech", err.getvalue())
 
 class AttributeErrorTests(unittest.TestCase):
     def test_attributes(self):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2021-04-14-03-53-06.bpo-38530.rNI_G1.rst b/Misc/NEWS.d/next/Core and Builtins/2021-04-14-03-53-06.bpo-38530.rNI_G1.rst
new file mode 100644
index 0000000000..ca175e7beb
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2021-04-14-03-53-06.bpo-38530.rNI_G1.rst
@@ -0,0 +1,3 @@
+When printing :exc:`NameError` raised by the interpreter,
+:c:func:`PyErr_Display` will offer suggestions of similar variable names in
+the function that the exception was raised from. Patch by Pablo Galindo
diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 4bb4153311..9916ce8854 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -1326,8 +1326,69 @@ SimpleExtendsException(PyExc_RuntimeError, NotImplementedError,
 /*
  *    NameError extends Exception
  */
-SimpleExtendsException(PyExc_Exception, NameError,
-                       "Name not found globally.");
+
+static int
+NameError_init(PyNameErrorObject *self, PyObject *args, PyObject *kwds)
+{
+    static char *kwlist[] = {"name", NULL};
+    PyObject *name = NULL;
+
+    if (BaseException_init((PyBaseExceptionObject *)self, args, NULL) == -1) {
+        return -1;
+    }
+
+    PyObject *empty_tuple = PyTuple_New(0);
+    if (!empty_tuple) {
+        return -1;
+    }
+    if (!PyArg_ParseTupleAndKeywords(empty_tuple, kwds, "|$O:NameError", kwlist,
+                                     &name)) {
+        Py_DECREF(empty_tuple);
+        return -1;
+    }
+    Py_DECREF(empty_tuple);
+
+    Py_XINCREF(name);
+    Py_XSETREF(self->name, name);
+
+    return 0;
+}
+
+static int
+NameError_clear(PyNameErrorObject *self)
+{
+    Py_CLEAR(self->name);
+    return BaseException_clear((PyBaseExceptionObject *)self);
+}
+
+static void
+NameError_dealloc(PyNameErrorObject *self)
+{
+    _PyObject_GC_UNTRACK(self);
+    NameError_clear(self);
+    Py_TYPE(self)->tp_free((PyObject *)self);
+}
+
+static int
+NameError_traverse(PyNameErrorObject *self, visitproc visit, void *arg)
+{
+    Py_VISIT(self->name);
+    return BaseException_traverse((PyBaseExceptionObject *)self, visit, arg);
+}
+
+static PyMemberDef NameError_members[] = {
+        {"name", T_OBJECT, offsetof(PyNameErrorObject, name), 0, PyDoc_STR("name")},
+        {NULL}  /* Sentinel */
+};
+
+static PyMethodDef NameError_methods[] = {
+        {NULL}  /* Sentinel */
+};
+
+ComplexExtendsException(PyExc_Exception, NameError,
+                        NameError, 0,
+                        NameError_methods, NameError_members,
+                        0, BaseException_str, "Name not found globally.");
 
 /*
  *    UnboundLocalError extends NameError
diff --git a/Python/ceval.c b/Python/ceval.c
index 53b596b304..326930b706 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -6319,6 +6319,20 @@ format_exc_check_arg(PyThreadState *tstate, PyObject *exc,
         return;
 
     _PyErr_Format(tstate, exc, format_str, obj_str);
+
+    if (exc == PyExc_NameError) {
+        // Include the name in the NameError exceptions to offer suggestions later.
+        _Py_IDENTIFIER(name);
+        PyObject *type, *value, *traceback;
+        PyErr_Fetch(&type, &value, &traceback);
+        PyErr_NormalizeException(&type, &value, &traceback);
+        if (PyErr_GivenExceptionMatches(value, PyExc_NameError)) {
+            // We do not care if this fails because we are going to restore the
+            // NameError anyway.
+            (void)_PyObject_SetAttrId(value, &PyId_name, obj);
+        }
+        PyErr_Restore(type, value, traceback);
+    }
 }
 
 static void
diff --git a/Python/suggestions.c b/Python/suggestions.c
index 2c0858d558..058294fc8b 100644
--- a/Python/suggestions.c
+++ b/Python/suggestions.c
@@ -1,17 +1,15 @@
 #include "Python.h"
+#include "frameobject.h"
 
 #include "pycore_pyerrors.h"
 
 #define MAX_DISTANCE 3
 #define MAX_CANDIDATE_ITEMS 100
-#define MAX_STRING_SIZE 20
+#define MAX_STRING_SIZE 25
 
 /* Calculate the Levenshtein distance between string1 and string2 */
 static size_t
 levenshtein_distance(const char *a, const char *b) {
-    if (a == NULL || b == NULL) {
-        return 0;
-    }
 
     const size_t a_size = strlen(a);
     const size_t b_size = strlen(b);
@@ -89,14 +87,19 @@ calculate_suggestions(PyObject *dir,
 
     Py_ssize_t suggestion_distance = PyUnicode_GetLength(name);
     PyObject *suggestion = NULL;
+    const char *name_str = PyUnicode_AsUTF8(name);
+    if (name_str == NULL) {
+        PyErr_Clear();
+        return NULL;
+    }
     for (int i = 0; i < dir_size; ++i) {
         PyObject *item = PyList_GET_ITEM(dir, i);
-        const char *name_str = PyUnicode_AsUTF8(name);
-        if (name_str == NULL) {
+        const char *item_str = PyUnicode_AsUTF8(item);
+        if (item_str == NULL) {
             PyErr_Clear();
-            continue;
+            return NULL;
         }
-        Py_ssize_t current_distance = levenshtein_distance(PyUnicode_AsUTF8(name), PyUnicode_AsUTF8(item));
+        Py_ssize_t current_distance = levenshtein_distance(name_str, item_str);
         if (current_distance == 0 || current_distance > MAX_DISTANCE) {
             continue;
         }
@@ -132,6 +135,48 @@ offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc) {
     return suggestions;
 }
 
+
+static PyObject *
+offer_suggestions_for_name_error(PyNameErrorObject *exc) {
+    PyObject *name = exc->name; // borrowed reference
+    PyTracebackObject *traceback = (PyTracebackObject *) exc->traceback; // borrowed reference
+    // Abort if we don't have an attribute name or we have an invalid one
+    if (name == NULL || traceback == NULL || !PyUnicode_CheckExact(name)) {
+        return NULL;
+    }
+
+    // Move to the traceback of the exception
+    while (traceback->tb_next != NULL) {
+        traceback = traceback->tb_next;
+    }
+
+    PyFrameObject *frame = traceback->tb_frame;
+    assert(frame != NULL);
+    PyCodeObject *code = frame->f_code;
+    assert(code != NULL && code->co_varnames != NULL);
+    PyObject *dir = PySequence_List(code->co_varnames);
+    if (dir == NULL) {
+        PyErr_Clear();
+        return NULL;
+    }
+
+    PyObject *suggestions = calculate_suggestions(dir, name);
+    Py_DECREF(dir);
+    if (suggestions != NULL) {
+        return suggestions;
+    }
+
+    dir = PySequence_List(frame->f_globals);
+    if (dir == NULL) {
+        PyErr_Clear();
+        return NULL;
+    }
+    suggestions = calculate_suggestions(dir, name);
+    Py_DECREF(dir);
+
+    return suggestions;
+}
+
 // Offer suggestions for a given exception. Returns a python string object containing the
 // suggestions. This function does not raise exceptions and returns NULL if no suggestion was found.
 PyObject *_Py_Offer_Suggestions(PyObject *exception) {
@@ -139,6 +184,8 @@ PyObject *_Py_Offer_Suggestions(PyObject *exception) {
     assert(!PyErr_Occurred()); // Check that we are not going to clean any existing exception
     if (PyErr_GivenExceptionMatches(exception, PyExc_AttributeError)) {
         result = offer_suggestions_for_attribute_error((PyAttributeErrorObject *) exception);
+    } else if (PyErr_GivenExceptionMatches(exception, PyExc_NameError)) {
+        result = offer_suggestions_for_name_error((PyNameErrorObject *) exception);
     }
     assert(!PyErr_Occurred());
     return result;