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)
# 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 = [
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)
@ -1168,6 +1174,23 @@ class ANSIString(str, metaclass=ANSIMeta):
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):
"""
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.lower().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):
"""
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")
first = ("\x1b[1m\x1b[32mThis is \x1b[0m", "This is ")
second = ("\x1b[1m\x1b[32m\x1b[0m split string\x1b[1m\x1b[32m", " split string")
target = ANSIString("|gThis is \nA split string|g")
first = ("\x1b[1m\x1b[32mThis is \n", "This is \n")
# 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)
normal_split = target.split("A")
self.assertEqual(re_split, normal_split)
@ -219,19 +220,19 @@ class ANSIStringTestCase(TestCase):
def test_slice_insert_longer(self):
"""
The ANSIString replays the color code before the split in order to
produce a *visually* identical result. The result is a longer string in
raw characters, but one which correctly represents the color output.
Test that slicing and inserting produces correct ANSI output,
with color codes preserved across newlines.
"""
string = ANSIString("A bigger |rTest|n of things |bwith more color|n")
# from evennia import set_trace;set_trace()
string = ANSIString("A bigger |rTest\n of things |bwith more color|n")
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(
repr(
(
ANSIString("A bigger ")
+ ANSIString("|rTest") # note that the |r|n is replayed together on next line
+ ANSIString("|r|n of things |bwith more color|n")
+ ANSIString("|rTest")
+ ANSIString("|r\n of things |bwith more color|n")
).raw()
),
repr(split_string.raw()),