Skip to content

Commit f0aeabc

Browse files
gh-140049: Colorize exception notes in traceback.py (#140051)
Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
1 parent 3776aba commit f0aeabc

File tree

4 files changed

+34
-6
lines changed

4 files changed

+34
-6
lines changed

Lib/_colorize.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,7 @@ class Syntax(ThemeSection):
327327
class Traceback(ThemeSection):
328328
type: str = ANSIColors.BOLD_MAGENTA
329329
message: str = ANSIColors.MAGENTA
330+
note: str = ANSIColors.CYAN
330331
filename: str = ANSIColors.MAGENTA
331332
line_no: str = ANSIColors.MAGENTA
332333
frame: str = ANSIColors.MAGENTA

Lib/test/test_traceback.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@
4040
test_frame = namedtuple('frame', ['f_code', 'f_globals', 'f_locals'])
4141
test_tb = namedtuple('tb', ['tb_frame', 'tb_lineno', 'tb_next', 'tb_lasti'])
4242

43-
color_overrides = {"reset": "z", "filename": "fn", "error_highlight": "E"}
43+
color_overrides = {"reset": "z", "filename": "fn", "error_highlight": "E", "note": "n"}
4444
colors = {
4545
color_overrides.get(k, k[0].lower()): v
4646
for k, v in _colorize.default_theme.traceback.items()
@@ -5306,6 +5306,23 @@ def bar():
53065306
self.assertIn("return baz1(1,\n 2,3\n ,4)", lines)
53075307
self.assertIn(red + "bar" + reset + boldr + "()" + reset, lines)
53085308

5309+
def test_colorized_exception_notes(self):
5310+
def foo():
5311+
raise ValueError()
5312+
5313+
try:
5314+
foo()
5315+
except Exception as e:
5316+
e.add_note("First note")
5317+
e.add_note("Second note")
5318+
exc = traceback.TracebackException.from_exception(e)
5319+
5320+
lines = "".join(exc.format(colorize=True))
5321+
note = colors["n"]
5322+
reset = colors["z"]
5323+
self.assertIn(note + "First note" + reset, lines)
5324+
self.assertIn(note + "Second note" + reset, lines)
5325+
53095326
def test_colorized_syntax_error(self):
53105327
try:
53115328
compile("a $ b", "<string>", "exec")
@@ -5314,7 +5331,7 @@ def test_colorized_syntax_error(self):
53145331
e, capture_locals=True
53155332
)
53165333
actual = "".join(exc.format(colorize=True))
5317-
def expected(t, m, fn, l, f, E, e, z):
5334+
def expected(t, m, fn, l, f, E, e, z, n):
53185335
return "".join(
53195336
[
53205337
f' File {fn}"<string>"{z}, line {l}1{z}\n',
@@ -5340,7 +5357,7 @@ def foo():
53405357
actual = tbstderr.getvalue().splitlines()
53415358

53425359
lno_foo = foo.__code__.co_firstlineno
5343-
def expected(t, m, fn, l, f, E, e, z):
5360+
def expected(t, m, fn, l, f, E, e, z, n):
53445361
return [
53455362
'Traceback (most recent call last):',
53465363
f' File {fn}"{__file__}"{z}, '
@@ -5373,7 +5390,7 @@ def foo():
53735390

53745391
lno_foo = foo.__code__.co_firstlineno
53755392
actual = "".join(exc.format(colorize=True)).splitlines()
5376-
def expected(t, m, fn, l, f, E, e, z):
5393+
def expected(t, m, fn, l, f, E, e, z, n):
53775394
return [
53785395
f" + Exception Group Traceback (most recent call last):",
53795396
f' | File {fn}"{__file__}"{z}, line {l}{lno_foo+9}{z}, in {f}test_colorized_traceback_from_exception_group{z}',

Lib/traceback.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -993,6 +993,10 @@ def _display_width(line, offset=None):
993993
)
994994

995995

996+
def _format_note(note, indent, theme):
997+
for l in note.split("\n"):
998+
yield f"{indent}{theme.note}{l}{theme.reset}\n"
999+
9961000

9971001
class _ExceptionPrintContext:
9981002
def __init__(self):
@@ -1291,6 +1295,10 @@ def format_exception_only(self, *, show_group=False, _depth=0, **kwargs):
12911295
well, recursively, with indentation relative to their nesting depth.
12921296
"""
12931297
colorize = kwargs.get("colorize", False)
1298+
if colorize:
1299+
theme = _colorize.get_theme(force_color=True).traceback
1300+
else:
1301+
theme = _colorize.get_theme(force_no_color=True).traceback
12941302

12951303
indent = 3 * _depth * ' '
12961304
if not self._have_exc_type:
@@ -1319,9 +1327,10 @@ def format_exception_only(self, *, show_group=False, _depth=0, **kwargs):
13191327
):
13201328
for note in self.__notes__:
13211329
note = _safe_string(note, 'note')
1322-
yield from [indent + l + '\n' for l in note.split('\n')]
1330+
yield from _format_note(note, indent, theme)
13231331
elif self.__notes__ is not None:
1324-
yield indent + "{}\n".format(_safe_string(self.__notes__, '__notes__', func=repr))
1332+
note = _safe_string(self.__notes__, '__notes__', func=repr)
1333+
yield from _format_note(note, indent, theme)
13251334

13261335
if self.exceptions and show_group:
13271336
for ex in self.exceptions:
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
:func:`traceback.format_exception_only` now colorizes exception notes.

0 commit comments

Comments
 (0)