@@ -974,11 +974,10 @@ def style():
974974 return Layout (new_layout )
975975
976976def 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):
10171016class 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
11711193def keyboard_interrupt_handler (signum , frame ):
0 commit comments