Fix ANSI escape explosion when slicing ANSIString after reset

This commit is contained in:
Count Infinity 2025-11-26 02:28:24 -07:00
parent d6d3fd3424
commit d9def9bfea
3 changed files with 68 additions and 12 deletions

View file

@ -1078,8 +1078,14 @@ class ANSIString(str, metaclass=ANSIMeta):
code_indexes_set = set(self._code_indexes) code_indexes_set = set(self._code_indexes)
# Only collect codes after the last reset to avoid accumulating
# cancelled codes when slicing
start_pos = self._find_last_reset_before(char_pos)
result_chars = [ result_chars = [
self._raw_string[index] for index in range(0, char_pos + 1) if index in code_indexes_set self._raw_string[index]
for index in range(start_pos, char_pos + 1)
if index in code_indexes_set
] ]
result = "".join(result_chars) result = "".join(result_chars)
@ -1168,6 +1174,23 @@ class ANSIString(str, metaclass=ANSIMeta):
return code_indexes, char_indexes return code_indexes, char_indexes
def _find_last_reset_before(self, pos):
"""
Find the end position of the last ANSI reset sequence
that occurs before the given position.
Args:
pos (int): Position in _raw_string to search before.
Returns:
int: The index immediately after the last reset sequence,
or 0 if no reset was found before pos.
"""
reset_pos = self._raw_string.rfind(ANSI_NORMAL, 0, pos)
if reset_pos == -1:
return 0
return reset_pos + len(ANSI_NORMAL)
def _get_interleving(self, index): def _get_interleving(self, index):
""" """
Get the code characters from the given slice end to the next Get the code characters from the given slice end to the next

View file

@ -195,3 +195,35 @@ class TestANSIString(TestCase):
self.assertEqual(plain.upper().clean(), "HELLO WORLD") self.assertEqual(plain.upper().clean(), "HELLO WORLD")
self.assertEqual(plain.lower().clean(), "hello world") self.assertEqual(plain.lower().clean(), "hello world")
self.assertEqual(plain.capitalize().clean(), "Hello world") self.assertEqual(plain.capitalize().clean(), "Hello world")
def test_getitem_no_cancelled_codes_after_reset(self):
"""
Test that slicing after a reset does NOT inherit cancelled codes.
This prevents exponential ANSI code accumulation during split/slice
operations. Text after a reset (|n) should not carry forward the
color codes that were cancelled by that reset.
"""
# String with red text, reset, then plain text
text = AN("|rRed|n plain")
# Slice starting after the reset - should NOT have red codes
after_reset = text[4:] # " plain"
self.assertEqual(after_reset.clean(), "plain")
# The raw output should NOT contain red color code since it was reset
raw = after_reset.raw()
self.assertNotIn(ANSI_RED, raw, "Cancelled red code should not appear after reset")
self.assertNotIn(ANSI_HILITE, raw, "Cancelled hilite code should not appear after reset")
# More complex case: multiple colors with resets
multi = AN("|rRed|n |gGreen|n |bBlue|n end")
# Slice after all color codes are reset
end_slice = multi[-3:] # "end"
self.assertEqual(end_slice.clean(), "end")
# Should have no color codes since all were reset
end_raw = end_slice.raw()
self.assertNotIn(ANSI_RED, end_raw)
self.assertNotIn(ANSI_GREEN, end_raw)
self.assertNotIn(ANSI_BLUE, end_raw)

View file

@ -94,11 +94,12 @@ class ANSIStringTestCase(TestCase):
def test_split(self): def test_split(self):
""" """
Verifies that re.split and .split behave similarly and that color Verifies that re.split and .split behave similarly and that color
codes end up where they should. codes end up where they should, including across newlines.
""" """
target = ANSIString("|gThis is |nA split string|g") target = ANSIString("|gThis is \nA split string|g")
first = ("\x1b[1m\x1b[32mThis is \x1b[0m", "This is ") first = ("\x1b[1m\x1b[32mThis is \n", "This is \n")
second = ("\x1b[1m\x1b[32m\x1b[0m split string\x1b[1m\x1b[32m", " split string") # Color codes carry through the newline
second = ("\x1b[1m\x1b[32m split string\x1b[1m\x1b[32m", " split string")
re_split = re.split("A", target) re_split = re.split("A", target)
normal_split = target.split("A") normal_split = target.split("A")
self.assertEqual(re_split, normal_split) self.assertEqual(re_split, normal_split)
@ -219,19 +220,19 @@ class ANSIStringTestCase(TestCase):
def test_slice_insert_longer(self): def test_slice_insert_longer(self):
""" """
The ANSIString replays the color code before the split in order to Test that slicing and inserting produces correct ANSI output,
produce a *visually* identical result. The result is a longer string in with color codes preserved across newlines.
raw characters, but one which correctly represents the color output.
""" """
string = ANSIString("A bigger |rTest|n of things |bwith more color|n") string = ANSIString("A bigger |rTest\n of things |bwith more color|n")
# from evennia import set_trace;set_trace()
split_string = string[:9] + "Test" + string[13:] split_string = string[:9] + "Test" + string[13:]
# string[:9] includes trailing red code since position 9 is in red region
# string[13:] carries red through the newline
self.assertEqual( self.assertEqual(
repr( repr(
( (
ANSIString("A bigger ") ANSIString("A bigger ")
+ ANSIString("|rTest") # note that the |r|n is replayed together on next line + ANSIString("|rTest")
+ ANSIString("|r|n of things |bwith more color|n") + ANSIString("|r\n of things |bwith more color|n")
).raw() ).raw()
), ),
repr(split_string.raw()), repr(split_string.raw()),