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

3 files changed, 24 insertions, 1 deletions

diff --git a/Lib/idlelib/idle_test/test_run.py b/Lib/idlelib/idle_test/test_run.py
index ec4637c..a38e43d 100644
--- a/Lib/idlelib/idle_test/test_run.py
+++ b/Lib/idlelib/idle_test/test_run.py
@@ -39,7 +39,8 @@ class ExceptionTest(unittest.TestCase):
 
     data = (('1/0', ZeroDivisionError, "division by zero\n"),
             ('abc', NameError, "name 'abc' is not defined. "
-                               "Did you mean: 'abs'?\n"),
+                               "Did you mean: 'abs'? "
+                               "Or did you forget to import 'abc'?\n"),
             ('int.reel', AttributeError,
                  "type object 'int' has no attribute 'reel'. "
                  "Did you mean: 'real'?\n"),
diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
index 4864b5c..2d17e06 100644
--- a/Lib/test/test_traceback.py
+++ b/Lib/test/test_traceback.py
@@ -3185,6 +3185,21 @@ class SuggestionFormattingTestBase:
         actual = self.get_suggestion(func)
         self.assertNotIn("something", actual)
 
+    def test_name_error_for_stdlib_modules(self):
+        def func():
+            stream = io.StringIO()
+
+        actual = self.get_suggestion(func)
+        self.assertIn("forget to import 'io'", actual)
+
+    def test_name_error_for_private_stdlib_modules(self):
+        def func():
+            stream = _io.StringIO()
+
+        actual = self.get_suggestion(func)
+        self.assertIn("forget to import '_io'", actual)
+
+
 
 class PurePythonSuggestionFormattingTests(
     PurePythonExceptionFormattingMixin,
diff --git a/Lib/traceback.py b/Lib/traceback.py
index c46ddaf..bb7856a 100644
--- a/Lib/traceback.py
+++ b/Lib/traceback.py
@@ -712,6 +712,13 @@ class TracebackException:
             suggestion = _compute_suggestion_error(exc_value, exc_traceback)
             if suggestion:
                 self._str += f". Did you mean: '{suggestion}'?"
+            if issubclass(exc_type, NameError):
+                wrong_name = getattr(exc_value, "name", None)
+                if wrong_name is not None and wrong_name in sys.stdlib_module_names:
+                    if suggestion:
+                        self._str += f" Or did you forget to import '{wrong_name}'"
+                    else:
+                        self._str += f". Did you forget to import '{wrong_name}'"
         if lookup_lines:
             self._load_lines()
         self.__suppress_context__ = \