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?

2 files changed, 69 insertions, 8 deletions

diff --git a/Python/ceval.c b/Python/ceval.c
index 53b596b..326930b 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 2c0858d..058294f 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;