Add more colors

pablogsal · pablogsal · commit 949da497b203 · 2023-12-05T14:13:48.000Z
diff --git a/Lib/test/test_traceback.py b/Lib/test/test_traceback.py
@@ -4300,18 +4300,19 @@ def foo():
 
         red = traceback._ANSIColors.RED
         boldr = traceback._ANSIColors.BOLD_RED
+        magenta = traceback._ANSIColors.MAGENTA
         reset = traceback._ANSIColors.RESET
         lno_foo = foo.__code__.co_firstlineno
         expected = ['Traceback (most recent call last):',
-            f'  File "{__file__}", '
-            f'line {lno_foo+5}, in test_colorized_traceback_is_the_default',
+            f'  File {magenta}"{__file__}"{reset}, '
+            f'line {magenta}{lno_foo+5}{reset}, in test_colorized_traceback_is_the_default',
             f'    {red}foo{reset+boldr}(){reset}',
             f'    {red}~~~{reset+boldr}^^{reset}',
-            f'  File "{__file__}", '
-            f'line {lno_foo+1}, in foo',
+            f'  File {magenta}"{__file__}"{reset}, '
+            f'line {magenta}{lno_foo+1}{reset}, in foo',
             f'    {red}1{reset+boldr}/{reset+red}0{reset}',
             f'    {red}~{reset+boldr}^{reset+red}~{reset}',
-            'ZeroDivisionError: division by zero']
+            f'{red}ZeroDivisionError{reset}: division by zero']
         self.assertEqual(actual, expected)
 
     def test_colorized_detection_checks_for_environment_variables(self):
diff --git a/Lib/traceback.py b/Lib/traceback.py
@@ -203,13 +203,19 @@ def format_exception_only(exc, /, value=_sentinel, *, show_group=False):
 
 # -- not official API but folk probably use these two functions.
 
-def _format_final_exc_line(etype, value, *, insert_final_newline=True):
+def _format_final_exc_line(etype, value, *, insert_final_newline=True, colorize=False):
     valuestr = _safe_string(value, 'exception')
     end_char = "\n" if insert_final_newline else ""
-    if value is None or not valuestr:
-        line = f"{etype}{end_char}"
+    if colorize:
+        if value is None or not valuestr:
+            line = f"{_ANSIColors.RED}{etype}{_ANSIColors.RESET}{end_char}"
+        else:
+            line = f"{_ANSIColors.RED}{etype}{_ANSIColors.RESET}: {valuestr}{end_char}"
     else:
-        line = f"{etype}: {valuestr}{end_char}"
+        if value is None or not valuestr:
+            line = f"{etype}{end_char}"
+        else:
+            line = f"{etype}: {valuestr}{end_char}"
     return line
 
 def _safe_string(value, what, func=str):
@@ -440,6 +446,8 @@ def _get_code_position(code, instruction_index):
 class _ANSIColors:
     RED = '\x1b[31m'
     BOLD_RED = '\x1b[1;31m'
+    MAGENTA = '\x1b[35m'
+    GREY = '\x1b[90m'
     RESET = '\x1b[0m'
 
 class StackSummary(list):
@@ -543,8 +551,20 @@ def format_frame_summary(self, frame_summary, **kwargs):
         filename = frame_summary.filename
         if frame_summary.filename.startswith("<stdin>-"):
             filename = "<stdin>"
-        row.append('  File "{}", line {}, in {}\n'.format(
-            filename, frame_summary.lineno, frame_summary.name))
+        if colorize:
+            row.append('  File {}"{}"{}, line {}{}{}, in {}\n'.format(
+                    _ANSIColors.MAGENTA,
+                    filename,
+                    _ANSIColors.RESET,
+                    _ANSIColors.MAGENTA,
+                    frame_summary.lineno,
+                    _ANSIColors.RESET,
+                    frame_summary.name,
+                    )
+            )
+        else:
+            row.append('  File "{}", line {}, in {}\n'.format(
+                filename, frame_summary.lineno, frame_summary.name))
         if frame_summary._dedented_lines and frame_summary._dedented_lines.strip():
             if (
                 frame_summary.colno is None or
@@ -1201,22 +1221,22 @@ def format_exception_only(self, *, show_group=False, _depth=0, **kwargs):
 
         indent = 3 * _depth * ' '
         if not self._have_exc_type:
-            yield indent + _format_final_exc_line(None, self._str)
+            yield indent + _format_final_exc_line(None, self._str, colorize=colorize)
             return
 
         stype = self.exc_type_str
         if not self._is_syntax_error:
             if _depth > 0:
                 # Nested exceptions needs correct handling of multiline messages.
                 formatted = _format_final_exc_line(
-                    stype, self._str, insert_final_newline=False,
+                    stype, self._str, insert_final_newline=False, colorize=colorize
                 ).split('\n')
                 yield from [
                     indent + l + '\n'
                     for l in formatted
                 ]
             else:
-                yield _format_final_exc_line(stype, self._str)
+                yield _format_final_exc_line(stype, self._str, colorize=colorize)
         else:
             yield from [indent + l for l in self._format_syntax_error(stype, colorize=colorize)]
 
@@ -1240,8 +1260,18 @@ def _format_syntax_error(self, stype, **kwargs):
         colorize = kwargs.get("colorize", False)
         filename_suffix = ''
         if self.lineno is not None:
-            yield '  File "{}", line {}\n'.format(
-                self.filename or "<string>", self.lineno)
+            if colorize:
+                yield '  File {}"{}"{}, line {}{}{}\n'.format(
+                    _ANSIColors.MAGENTA,
+                    self.filename or "<string>",
+                    _ANSIColors.RESET,
+                    _ANSIColors.MAGENTA,
+                    self.lineno,
+                    _ANSIColors.RESET,
+                    )
+            else:
+                yield '  File "{}", line {}\n'.format(
+                    self.filename or "<string>", self.lineno)
         elif self.filename is not None:
             filename_suffix = ' ({})'.format(self.filename)
 
