Skip to content

Commit

Permalink
Add more colors
Browse files Browse the repository at this point in the history
  • Loading branch information
pablogsal committed Dec 5, 2023
1 parent 9ae4c77 commit 949da49
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 16 deletions.
11 changes: 6 additions & 5 deletions Lib/test/test_traceback.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
52 changes: 41 additions & 11 deletions Lib/traceback.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)]

Expand All @@ -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)

Expand Down

0 comments on commit 949da49

Please sign in to comment.