gh-98254: Include stdlib module names in error messages for NameErrors (#98255)

2 files changed, 74 insertions, 30 deletions

diff --git a/Python/pythonrun.c b/Python/pythonrun.c
index acb1330..a0005b3 100644
--- a/Python/pythonrun.c
+++ b/Python/pythonrun.c
@@ -1107,16 +1107,9 @@ print_exception_suggestions(struct exception_print_context *ctx,
     PyObject *f = ctx->file;
     PyObject *suggestions = _Py_Offer_Suggestions(value);
     if (suggestions) {
-        // Add a trailer ". Did you mean: (...)?"
-        if (PyFile_WriteString(". Did you mean: '", f) < 0) {
-            goto error;
-        }
         if (PyFile_WriteObject(suggestions, f, Py_PRINT_RAW) < 0) {
             goto error;
         }
-        if (PyFile_WriteString("'?", f) < 0) {
-            goto error;
-        }
         Py_DECREF(suggestions);
     }
     else if (PyErr_Occurred()) {
diff --git a/Python/suggestions.c b/Python/suggestions.c
index c336ec8..89b86f7 100644
--- a/Python/suggestions.c
+++ b/Python/suggestions.c
@@ -3,6 +3,7 @@
 
 #include "pycore_pyerrors.h"
 #include "pycore_code.h"        // _PyCode_GetVarnames()
+#include "stdlib_module_names.h"  // _Py_stdlib_module_names
 
 #define MAX_CANDIDATE_ITEMS 750
 #define MAX_STRING_SIZE 40
@@ -175,7 +176,7 @@ calculate_suggestions(PyObject *dir,
 }
 
 static PyObject *
-offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
+get_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
 {
     PyObject *name = exc->name; // borrowed reference
     PyObject *obj = exc->obj; // borrowed reference
@@ -195,35 +196,25 @@ offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
     return suggestions;
 }
 
-
 static PyObject *
-offer_suggestions_for_name_error(PyNameErrorObject *exc)
+offer_suggestions_for_attribute_error(PyAttributeErrorObject *exc)
 {
-    PyObject *name = exc->name; // borrowed reference
-    PyTracebackObject *traceback = (PyTracebackObject *) exc->traceback; // borrowed reference
-    // Abort if we don't have a variable name or we have an invalid one
-    // or if we don't have a traceback to work with
-    if (name == NULL || !PyUnicode_CheckExact(name) ||
-        traceback == NULL || !Py_IS_TYPE(traceback, &PyTraceBack_Type)
-    ) {
+    PyObject* suggestion = get_suggestions_for_attribute_error(exc);
+    if (suggestion == NULL) {
         return NULL;
     }
+    // Add a trailer ". Did you mean: (...)?"
+    PyObject* result = PyUnicode_FromFormat(". Did you mean: %R?", suggestion);
+    Py_DECREF(suggestion);
+    return result;
+}
 
-    // Move to the traceback of the exception
-    while (1) {
-        PyTracebackObject *next = traceback->tb_next;
-        if (next == NULL || !Py_IS_TYPE(next, &PyTraceBack_Type)) {
-            break;
-        }
-        else {
-            traceback = next;
-        }
-    }
-
-    PyFrameObject *frame = traceback->tb_frame;
-    assert(frame != NULL);
+static PyObject *
+get_suggestions_for_name_error(PyObject* name, PyFrameObject* frame)
+{
     PyCodeObject *code = PyFrame_GetCode(frame);
     assert(code != NULL && code->co_localsplusnames != NULL);
+
     PyObject *varnames = _PyCode_GetVarnames(code);
     if (varnames == NULL) {
         return NULL;
@@ -261,6 +252,66 @@ offer_suggestions_for_name_error(PyNameErrorObject *exc)
     return suggestions;
 }
 
+static bool
+is_name_stdlib_module(PyObject* name)
+{
+    const char* the_name = PyUnicode_AsUTF8(name);
+    Py_ssize_t len = Py_ARRAY_LENGTH(_Py_stdlib_module_names);
+    for (Py_ssize_t i = 0; i < len; i++) {
+        if (strcmp(the_name, _Py_stdlib_module_names[i]) == 0) {
+            return 1;
+        }
+    }
+    return 0;
+}
+
+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 a variable name or we have an invalid one
+    // or if we don't have a traceback to work with
+    if (name == NULL || !PyUnicode_CheckExact(name) ||
+        traceback == NULL || !Py_IS_TYPE(traceback, &PyTraceBack_Type)
+    ) {
+        return NULL;
+    }
+
+    // Move to the traceback of the exception
+    while (1) {
+        PyTracebackObject *next = traceback->tb_next;
+        if (next == NULL || !Py_IS_TYPE(next, &PyTraceBack_Type)) {
+            break;
+        }
+        else {
+            traceback = next;
+        }
+    }
+
+    PyFrameObject *frame = traceback->tb_frame;
+    assert(frame != NULL);
+
+    PyObject* suggestion = get_suggestions_for_name_error(name, frame);
+    bool is_stdlib_module = is_name_stdlib_module(name);
+
+    if (suggestion == NULL && !is_stdlib_module) {
+        return NULL;
+    }
+
+    // Add a trailer ". Did you mean: (...)?"
+    PyObject* result = NULL;
+    if (!is_stdlib_module) {
+        result = PyUnicode_FromFormat(". Did you mean: %R?", suggestion);
+    } else if (suggestion == NULL) {
+        result = PyUnicode_FromFormat(". Did you forget to import %R?", name);
+    } else {
+        result = PyUnicode_FromFormat(". Did you mean: %R? Or did you forget to import %R?", suggestion, name);
+    }
+    Py_XDECREF(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.