gh-99139: Improve NameError error suggestion for instances (#99140)

2 files changed, 35 insertions, 0 deletions

diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
index 149d023..430daf6 100644
--- a/Lib/test/test_traceback.py
+++ b/Lib/test/test_traceback.py
@@ -3356,6 +3356,31 @@ class SuggestionFormattingTestBase:
 
         actual = self.get_suggestion(func)
         self.assertNotIn("blech", actual)
+    
+    def test_name_error_with_instance(self):
+        class A:
+            def __init__(self):
+                self.blech = None
+            def foo(self):
+                blich = 1
+                x = blech
+
+        instance = A()
+        actual = self.get_suggestion(instance.foo)
+        self.assertIn("self.blech", actual)
+
+    def test_unbound_local_error_with_instance(self):
+        class A:
+            def __init__(self):
+                self.blech = None
+            def foo(self):
+                blich = 1
+                x = blech
+                blech = 1
+
+        instance = A()
+        actual = self.get_suggestion(instance.foo)
+        self.assertNotIn("self.blech", actual)
 
     def test_unbound_local_error_does_not_match(self):
         def func():
diff --git a/Lib/traceback.py b/Lib/traceback.py
index cf5f355..8d51872 100644
--- a/Lib/traceback.py
+++ b/Lib/traceback.py
@@ -1037,6 +1037,16 @@ def _compute_suggestion_error(exc_value, tb, wrong_name):
             + list(frame.f_globals)
             + list(frame.f_builtins)
         )
+
+        # Check first if we are in a method and the instance
+        # has the wrong name as attribute
+        if 'self' in frame.f_locals:
+            self = frame.f_locals['self']
+            if hasattr(self, wrong_name):
+                return f"self.{wrong_name}"
+    
+    # Compute closest match
+
     if len(d) > _MAX_CANDIDATE_ITEMS:
         return None
     wrong_name_len = len(wrong_name)