fix(traceback): TracebackException to handle error messages with punctuations

Author: krisztian-gajdar

Lib/traceback.py

@@ -1081,50 +1081,30 @@ def __init__(self, exc_type, exc_value, exc_traceback, *, limit=None,
 
         self._is_syntax_error = False
         self._have_exc_type = exc_type is not None
-        if exc_type is not None:
+            
+        if self._have_exc_type:
             self.exc_type_qualname = exc_type.__qualname__
             self.exc_type_module = exc_type.__module__
+            if issubclass(exc_type, SyntaxError):
+                # Handle SyntaxError's specially
+                self.filename = exc_value.filename
+                lno = exc_value.lineno
+                self.lineno = str(lno) if lno is not None else None
+                end_lno = exc_value.end_lineno
+                self.end_lineno = str(end_lno) if end_lno is not None else None
+                self.text = exc_value.text
+                self.offset = exc_value.offset
+                self.end_offset = exc_value.end_offset
+                self.msg = exc_value.msg
+                self._is_syntax_error = True
+                self._exc_metadata = getattr(exc_value, "_metadata", None)
+            elif suggestion := _compute_suggestion_message(exc_type, exc_value, exc_traceback):
+                punctuation = "" if self._str.rstrip().endswith(('.', '!')) else "."
+                self._str += f"{punctuation} {suggestion}"
         else:
             self.exc_type_qualname = None
             self.exc_type_module = None
-
-        if exc_type and issubclass(exc_type, SyntaxError):
-            # Handle SyntaxError's specially
-            self.filename = exc_value.filename
-            lno = exc_value.lineno
-            self.lineno = str(lno) if lno is not None else None
-            end_lno = exc_value.end_lineno
-            self.end_lineno = str(end_lno) if end_lno is not None else None
-            self.text = exc_value.text
-            self.offset = exc_value.offset
-            self.end_offset = exc_value.end_offset
-            self.msg = exc_value.msg
-            self._is_syntax_error = True
-            self._exc_metadata = getattr(exc_value, "_metadata", None)
-        elif exc_type and issubclass(exc_type, ImportError) and \
-                getattr(exc_value, "name_from", None) is not None:
-            wrong_name = getattr(exc_value, "name_from", None)
-            suggestion = _compute_suggestion_error(exc_value, exc_traceback, wrong_name)
-            if suggestion:
-                self._str += f". Did you mean: '{suggestion}'?"
-        elif exc_type and issubclass(exc_type, ModuleNotFoundError) and \
-                sys.flags.no_site and \
-                getattr(exc_value, "name", None) not in sys.stdlib_module_names:
-            self._str += (". Site initialization is disabled, did you forget to "
-                + "add the site-packages directory to sys.path?")
-        elif exc_type and issubclass(exc_type, (NameError, AttributeError)) and \
-                getattr(exc_value, "name", None) is not None:
-            wrong_name = getattr(exc_value, "name", None)
-            suggestion = _compute_suggestion_error(exc_value, exc_traceback, wrong_name)
-            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__ = \
@@ -1384,7 +1364,6 @@ def _find_keyword_typos(self):
                 self.msg = f"invalid syntax. Did you mean '{suggestion}'?"
                 return
 
-
     def _format_syntax_error(self, stype, **kwargs):
         """Format SyntaxError exceptions (internal helper)."""
         # Show exactly where the problem was found.
@@ -1577,7 +1556,6 @@ def format(self, *, chain=True, _ctx=None, **kwargs):
                     assert _ctx.exception_group_depth == 1
                     _ctx.exception_group_depth = 0
 
-
     def print(self, *, file=None, chain=True, **kwargs):
         """Print the result of self.format(chain=chain) to 'file'."""
         colorize = kwargs.get("colorize", False)
@@ -1732,6 +1710,26 @@ def _compute_suggestion_error(exc_value, tb, wrong_name):
     return suggestion
 
 
+def _compute_suggestion_message(exc_type, exc_value, exc_traceback) -> str | None:
+    if issubclass(exc_type, ModuleNotFoundError) and \
+        sys.flags.no_site and \
+        getattr(exc_value, "name", None) not in sys.stdlib_module_names:
+            return ("Site initialization is disabled, did you forget to " 
+                "add the site-packages directory to sys.path?")
+    if issubclass(exc_type, (ImportError, NameError, AttributeError)):
+        name_attribute = "name_from" if issubclass(exc_type, ImportError) else "name"
+        wrong_name = getattr(exc_value, name_attribute, None)
+        if wrong_name:
+            suggestion_error = _compute_suggestion_error(exc_value, exc_traceback, wrong_name)
+            suggestion_text = f"Did you mean: '{suggestion_error}'?" if suggestion_error else None
+            if issubclass(exc_type, NameError) and wrong_name in sys.stdlib_module_names:
+                if suggestion_error:
+                    return f"{suggestion_text} Or did you forget to import '{wrong_name}'?"
+                return f"Did you forget to import '{wrong_name}'?"
+            return suggestion_text
+    return None
+
+
 def _levenshtein_distance(a, b, max_cost):
     # A Python implementation of Python/suggestions.c:levenshtein_distance.