Skip to content

Commit 2f4dd2c

Browse files
asmeurercodex
andcommitted
Fix traceback highlight edge cases
Handle wide-character spans, nested traceback exceptions, and chain ordering for inline traceback highlighting, and add regression tests for those cases. Co-authored-by: OpenAI Codex <codex@openai.com>
1 parent 7449ded commit 2f4dd2c

2 files changed

Lines changed: 207 additions & 77 deletions

File tree

‎mypython/mypython.py‎

Lines changed: 99 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -974,11 +974,10 @@ def style():
974974
return Layout(new_layout)
975975

976976
def add_ansi_at_columns(ansi_line, start_col, end_col, start_code, end_code):
977-
"""Insert ANSI escape codes at visible-column positions in an ANSI-formatted line.
977+
"""Insert ANSI escape codes at visible-character positions in an ANSI line.
978978
979-
Walks the line character by character, tracking visible column position
980-
(skipping ANSI escape sequences), and inserts start_code before start_col
981-
and end_code after end_col.
979+
start_col and end_col are 0-based offsets into the rendered source text
980+
after ANSI escape sequences are stripped.
982981
"""
983982
result = []
984983
visible_col = 0
@@ -1017,9 +1016,13 @@ def add_ansi_at_columns(ansi_line, start_col, end_col, start_code, end_code):
10171016
class MyStackSummary(traceback.StackSummary):
10181017
def __init__(self):
10191018
super().__init__()
1020-
# List of (stripped_line_text, start_visible_col, end_visible_col)
1019+
# List of (stripped_line_text, start_char_offset, end_char_offset)
10211020
self.highlight_ranges = []
10221021

1022+
def format(self):
1023+
self.highlight_ranges = []
1024+
return super().format()
1025+
10231026
def format_frame_summary(self, frame_summary):
10241027
"""Format a frame summary without caret lines, recording highlight ranges instead."""
10251028
row = []
@@ -1036,24 +1039,18 @@ def format_frame_summary(self, frame_summary):
10361039
and frame_summary.end_colno is not None
10371040
):
10381041
line = frame_summary._original_line
1039-
orig_line_len = len(line)
1040-
frame_line_len = len(frame_summary.line.lstrip())
1041-
stripped_characters = orig_line_len - frame_line_len
1042+
leading_whitespace = len(line) - len(line.lstrip())
10421043

10431044
start_offset = traceback._byte_offset_to_character_offset(
10441045
line, frame_summary.colno)
10451046
end_offset = traceback._byte_offset_to_character_offset(
10461047
line, frame_summary.end_colno)
10471048

10481049
if frame_summary.lineno != frame_summary.end_lineno:
1049-
end_offset = len(line.rstrip())
1050-
1051-
dp_start_offset = traceback._display_width(line, start_offset) + 1
1052-
dp_end_offset = traceback._display_width(line, end_offset) + 1
1050+
end_offset = len(line.rstrip('\n'))
10531051

1054-
# Record the highlight range (0-based visible columns within the stripped line)
1055-
highlight_start = dp_start_offset - stripped_characters
1056-
highlight_end = dp_end_offset - stripped_characters
1052+
highlight_start = start_offset - leading_whitespace
1053+
highlight_end = end_offset - leading_whitespace
10571054
if highlight_start >= 0 and highlight_end > highlight_start:
10581055
self.highlight_ranges.append((stripped_line, highlight_start, highlight_end))
10591056

@@ -1063,61 +1060,82 @@ def format_frame_summary(self, frame_summary):
10631060

10641061
return ''.join(row)
10651062

1066-
def _convert_chained_stacks(exc):
1067-
"""Walk the exception chain and convert all stacks to MyStackSummary."""
1068-
for attr in ('__cause__', '__context__'):
1069-
chained = getattr(exc, attr, None)
1070-
if chained is not None and not isinstance(chained.stack, MyStackSummary):
1063+
def _convert_traceback_stack(exc, remove_mypython):
1064+
"""Convert a traceback stack to MyStackSummary and optionally filter mypython frames."""
1065+
new_stack = MyStackSummary()
1066+
mypython_error = None
1067+
frame = None
1068+
for frame in exc.stack[:]:
1069+
if frame.filename.startswith(mypython_dir):
1070+
if mypython_error is False:
1071+
mypython_error = True
1072+
elif frame.filename.startswith('<mypython'):
1073+
if DOCTEST_MODE:
1074+
filename, lineno, name, line = frame
1075+
new_stack.append(traceback.FrameSummary("<stdin>", lineno, name))
1076+
else:
1077+
new_stack.append(frame)
1078+
else:
1079+
new_stack.append(frame)
1080+
if not mypython_error:
1081+
mypython_error = False
1082+
1083+
if (
1084+
frame is not None
1085+
and mypython_error is None
1086+
and frame.filename.startswith(mypython_dir)
1087+
and frame.name != 'keyboard_interrupt_handler'
1088+
):
1089+
mypython_error = True
1090+
1091+
if remove_mypython and not mypython_error:
1092+
exc.stack = new_stack
1093+
else:
1094+
if isinstance(exc.stack, MyStackSummary):
1095+
exc.stack.highlight_ranges = []
1096+
else:
10711097
my_stack = MyStackSummary()
1072-
my_stack.extend(chained.stack)
1073-
chained.stack = my_stack
1074-
_convert_chained_stacks(chained)
1098+
my_stack.extend(exc.stack)
1099+
exc.stack = my_stack
10751100

1076-
class MyTracebackException(traceback.TracebackException):
1077-
def __init__(self, exc_type, exc_value, exc_traceback, *,
1078-
remove_mypython=True, **kwargs):
1101+
exc.mypython_error = mypython_error
1102+
exc.syntax_highlight_ranges = []
10791103

1080-
super().__init__(exc_type, exc_value, exc_traceback, **kwargs)
1104+
def _rewrap_traceback_exceptions(exc, remove_mypython, seen=None):
1105+
"""Convert a TracebackException graph to MyTracebackException/MyStackSummary."""
1106+
if exc is None:
1107+
return
10811108

1082-
new_stack = MyStackSummary()
1083-
mypython_error = None
1084-
for frame in self.stack[:]:
1085-
if frame.filename.startswith(mypython_dir):
1086-
if mypython_error is False:
1087-
mypython_error = True
1088-
elif frame.filename.startswith('<mypython'):
1089-
if DOCTEST_MODE:
1090-
filename, lineno, name, line = frame
1091-
new_stack.append(traceback.FrameSummary("<stdin>", lineno, name))
1092-
else:
1093-
new_stack.append(frame)
1094-
else:
1095-
new_stack.append(frame)
1096-
if not mypython_error:
1097-
mypython_error = False
1109+
if seen is None:
1110+
seen = set()
10981111

1099-
if (mypython_error is None and frame.filename.startswith(mypython_dir)
1100-
and frame.name != 'keyboard_interrupt_handler'):
1101-
mypython_error = True
1112+
exc_id = id(exc)
1113+
if exc_id in seen:
1114+
return
1115+
seen.add(exc_id)
11021116

1103-
if remove_mypython and not mypython_error:
1104-
self.stack = new_stack
1105-
else:
1106-
# Still convert to MyStackSummary for highlight support
1107-
my_stack = MyStackSummary()
1108-
my_stack.extend(self.stack)
1109-
self.stack = my_stack
1117+
if not isinstance(exc, MyTracebackException):
1118+
exc.__class__ = MyTracebackException
11101119

1111-
self.mypython_error = mypython_error
1120+
_convert_traceback_stack(exc, remove_mypython)
11121121

1113-
# Also convert stacks on chained exceptions
1114-
_convert_chained_stacks(self)
1122+
for attr in ('__cause__', '__context__'):
1123+
_rewrap_traceback_exceptions(getattr(exc, attr, None), remove_mypython, seen)
11151124

1116-
# Collect all syntax error highlight ranges
1117-
self.syntax_highlight_ranges = []
1125+
if exc.exceptions is not None:
1126+
for child in exc.exceptions:
1127+
_rewrap_traceback_exceptions(child, remove_mypython, seen)
1128+
1129+
class MyTracebackException(traceback.TracebackException):
1130+
def __init__(self, exc_type, exc_value, exc_traceback, *,
1131+
remove_mypython=True, **kwargs):
1132+
1133+
super().__init__(exc_type, exc_value, exc_traceback, **kwargs)
1134+
_rewrap_traceback_exceptions(self, remove_mypython)
11181135

11191136
def _format_syntax_error(self, stype):
11201137
"""Format SyntaxError without caret lines, recording highlight ranges."""
1138+
self.syntax_highlight_ranges = []
11211139
filename_suffix = ''
11221140
if self.lineno is not None:
11231141
yield ' File "{}", line {}\n'.format(
@@ -1146,26 +1164,30 @@ def _format_syntax_error(self, stype):
11461164
msg = self.msg or "<no detail available>"
11471165
yield "{}: {}{}\n".format(stype, msg, filename_suffix)
11481166

1149-
def collect_all_highlight_ranges(self):
1150-
"""Collect all highlight ranges from this exception and its chain."""
1167+
def _collect_highlight_ranges(self, ranges, chain, seen):
1168+
exc_id = id(self)
1169+
if exc_id in seen:
1170+
return
1171+
seen.add(exc_id)
1172+
1173+
if chain:
1174+
if self.__cause__ is not None:
1175+
self.__cause__._collect_highlight_ranges(ranges, chain, seen)
1176+
elif self.__context__ is not None and not self.__suppress_context__:
1177+
self.__context__._collect_highlight_ranges(ranges, chain, seen)
1178+
1179+
if isinstance(self.stack, MyStackSummary):
1180+
ranges.extend(self.stack.highlight_ranges)
1181+
ranges.extend(self.syntax_highlight_ranges)
1182+
1183+
if self.exceptions is not None:
1184+
for child in self.exceptions:
1185+
child._collect_highlight_ranges(ranges, chain, seen)
1186+
1187+
def collect_all_highlight_ranges(self, chain=True):
1188+
"""Collect highlight ranges in the same order TracebackException.format() prints them."""
11511189
ranges = []
1152-
exc = self
1153-
seen = set()
1154-
while exc is not None:
1155-
if id(exc) in seen:
1156-
break
1157-
seen.add(id(exc))
1158-
if isinstance(exc.stack, MyStackSummary):
1159-
ranges.extend(exc.stack.highlight_ranges)
1160-
if hasattr(exc, 'syntax_highlight_ranges'):
1161-
ranges.extend(exc.syntax_highlight_ranges)
1162-
# Walk the chain
1163-
if exc.__cause__ is not None:
1164-
exc = exc.__cause__
1165-
elif exc.__context__ is not None and not exc.__suppress_context__:
1166-
exc = exc.__context__
1167-
else:
1168-
break
1190+
self._collect_highlight_ranges(ranges, chain, set())
11691191
return ranges
11701192

11711193
def keyboard_interrupt_handler(signum, frame):

‎mypython/tests/test_mypython.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -912,6 +912,30 @@ def test_apply_inline_highlights():
912912
result = _apply_inline_highlights(ansi, ranges)
913913
assert result.count(start) == 2
914914

915+
@pytest.mark.skipif(sys.version_info < (3, 11),
916+
reason="colno/end_colno requires Python 3.11+")
917+
def test_my_stack_summary_highlight_ranges_wide_unicode():
918+
import tempfile, os
919+
920+
tf = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
921+
tf.write('变量 = 1/0\n')
922+
tf.close()
923+
try:
924+
exec(compile(open(tf.name).read(), tf.name, 'exec'))
925+
except:
926+
tbexc = MyTracebackException(*sys.exc_info(), remove_mypython=False)
927+
''.join(tbexc.format(chain=True))
928+
929+
ranges = tbexc.collect_all_highlight_ranges()
930+
error_ranges = [(t, s, e) for t, s, e in ranges if t == '变量 = 1/0']
931+
assert error_ranges == [('变量 = 1/0', 5, 8)]
932+
933+
start, end = TRACEBACK_HIGHLIGHT_STYLES['background']
934+
result = _apply_inline_highlights(' 变量 = 1/0', error_ranges)
935+
assert result == f' 变量 = {start}1/0{end}'
936+
finally:
937+
os.unlink(tf.name)
938+
915939
@pytest.mark.skipif(sys.version_info < (3, 11),
916940
reason="colno/end_colno requires Python 3.11+")
917941
def test_my_stack_summary_highlight_ranges():
@@ -982,6 +1006,90 @@ def test_syntax_error_highlight_ranges():
9821006
assert start == 7
9831007
assert end == 8
9841008

1009+
def test_nested_syntax_error_uses_custom_formatting():
1010+
try:
1011+
try:
1012+
compile('x = 1 +\n', '<test>', 'exec')
1013+
except SyntaxError as exc:
1014+
raise RuntimeError("wrapped") from exc
1015+
except RuntimeError:
1016+
tbexc = MyTracebackException(*sys.exc_info(), remove_mypython=False)
1017+
tb_str = ''.join(tbexc.format(chain=True))
1018+
1019+
assert re.search(r'^\s*[~^]+\s*$', tb_str, re.MULTILINE) is None
1020+
1021+
ranges = tbexc.collect_all_highlight_ranges()
1022+
syntax_ranges = [(t, s, e) for t, s, e in ranges if t == 'x = 1 +']
1023+
assert syntax_ranges == [('x = 1 +', 7, 8)]
1024+
1025+
@pytest.mark.skipif(sys.version_info < (3, 11),
1026+
reason="ExceptionGroup requires Python 3.11+")
1027+
def test_exception_group_children_use_custom_formatting():
1028+
import tempfile, os
1029+
1030+
tf = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
1031+
tf.write('def boom():\n x = 1/0\n\n'
1032+
'try:\n boom()\n'
1033+
'except Exception as exc:\n raise ExceptionGroup("group", [exc]) from None\n')
1034+
tf.close()
1035+
try:
1036+
namespace = {}
1037+
exec(compile(open(tf.name).read(), tf.name, 'exec'), namespace, namespace)
1038+
except ExceptionGroup:
1039+
tbexc = MyTracebackException(*sys.exc_info(), remove_mypython=False)
1040+
tb_str = ''.join(tbexc.format(chain=True))
1041+
1042+
assert re.search(r'^\s*[~^]+\s*$', tb_str, re.MULTILINE) is None
1043+
1044+
ranges = tbexc.collect_all_highlight_ranges()
1045+
error_ranges = [(t, s, e) for t, s, e in ranges if t == 'x = 1/0']
1046+
assert error_ranges == [('x = 1/0', 4, 7)]
1047+
finally:
1048+
os.unlink(tf.name)
1049+
1050+
@pytest.mark.skipif(sys.version_info < (3, 11),
1051+
reason="colno/end_colno requires Python 3.11+")
1052+
def test_collect_all_highlight_ranges_follows_chain_format_order(monkeypatch):
1053+
import tempfile, os
1054+
1055+
tf = tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False)
1056+
tf.write(
1057+
'def inner():\n'
1058+
' x, y = [1], []\n'
1059+
' return x[0] + y[0]\n'
1060+
'\n'
1061+
'def outer():\n'
1062+
' try:\n'
1063+
' inner()\n'
1064+
' except IndexError:\n'
1065+
' x, y = [], [1]\n'
1066+
' return x[0] + y[0]\n'
1067+
'\n'
1068+
'outer()\n'
1069+
)
1070+
tf.close()
1071+
try:
1072+
namespace = {}
1073+
exec(compile(open(tf.name).read(), tf.name, 'exec'), namespace, namespace)
1074+
except IndexError:
1075+
tbexc = MyTracebackException(*sys.exc_info(), remove_mypython=False)
1076+
tb_str = ''.join(tbexc.format(chain=True))
1077+
1078+
monkeypatch.setitem(
1079+
mypython.TRACEBACK_HIGHLIGHT_STYLES,
1080+
mypython.TRACEBACK_HIGHLIGHT_STYLE,
1081+
('<S>', '<E>'),
1082+
)
1083+
highlighted = _apply_inline_highlights(tb_str, tbexc.collect_all_highlight_ranges())
1084+
matching_lines = [line for line in highlighted.splitlines()
1085+
if 'return' in line and 'x[0]' in line and 'y[0]' in line]
1086+
assert matching_lines == [
1087+
' return x[0] + <S>y[0]<E>',
1088+
' return <S>x[0]<E> + y[0]',
1089+
]
1090+
finally:
1091+
os.unlink(tf.name)
1092+
9851093
def test_traceback_no_carets(check_output):
9861094
out, err = check_output('1/0\n')
9871095
assert '~^~' not in err

0 commit comments

Comments
 (0)