diff --git a/evennia/utils/ansi.py b/evennia/utils/ansi.py index 38bea3fcb1..f7b7a16218 100644 --- a/evennia/utils/ansi.py +++ b/evennia/utils/ansi.py @@ -703,12 +703,17 @@ def _transform(func_name): def wrapped(self, *args, **kwargs): replacement_string = _query_super(func_name)(self, *args, **kwargs) + + # Convert to sets for O(1) membership testing + code_indexes_set = set(self._code_indexes) + char_indexes_set = set(self._char_indexes) + to_string = [] char_counter = 0 for index in range(0, len(self._raw_string)): - if index in self._code_indexes: + if index in code_indexes_set: to_string.append(self._raw_string[index]) - elif index in self._char_indexes: + elif index in char_indexes_set: to_string.append(replacement_string[char_counter]) char_counter += 1 return ANSIString( @@ -1028,10 +1033,12 @@ class ANSIString(str, metaclass=ANSIMeta): return ANSIString("") last_mark = slice_indexes[0] # Check between the slice intervals for escape sequences. + # Convert to set for O(1) membership testing + code_indexes_set = set(self._code_indexes) i = None for i in slice_indexes[1:]: for index in range(last_mark, i): - if index in self._code_indexes: + if index in code_indexes_set: string += self._raw_string[index] last_mark = i try: @@ -1065,15 +1072,18 @@ class ANSIString(str, metaclass=ANSIMeta): append_tail = self._get_interleving(item + 1) else: append_tail = "" - item = self._char_indexes[item] - clean = self._raw_string[item] - result = "" - # Get the character they're after, and replay all escape sequences - # previous to it. - for index in range(0, item + 1): - if index in self._code_indexes: - result += self._raw_string[index] + char_pos = self._char_indexes[item] + clean = self._raw_string[char_pos] + + code_indexes_set = set(self._code_indexes) + + result_chars = [ + self._raw_string[index] for index in range(0, char_pos + 1) if index in code_indexes_set + ] + + result = "".join(result_chars) + return ANSIString(result + clean + append_tail, decoded=True) def clean(self): @@ -1153,7 +1163,9 @@ class ANSIString(str, metaclass=ANSIMeta): # Plain string, no ANSI codes. return code_indexes, list(range(0, len(self._raw_string))) # all indexes not occupied by ansi codes are normal characters - char_indexes = [i for i in range(len(self._raw_string)) if i not in code_indexes] + code_indexes_set = set(code_indexes) + char_indexes = [i for i in range(len(self._raw_string)) if i not in code_indexes_set] + return code_indexes, char_indexes def _get_interleving(self, index): @@ -1166,12 +1178,17 @@ class ANSIString(str, metaclass=ANSIMeta): index = self._char_indexes[index - 1] except IndexError: return "" + + # Convert to sets for O(1) membership testing + char_indexes_set = set(self._char_indexes) + code_indexes_set = set(self._code_indexes) + s = "" while True: index += 1 - if index in self._char_indexes: + if index in char_indexes_set: break - elif index in self._code_indexes: + elif index in code_indexes_set: s += self._raw_string[index] else: break diff --git a/evennia/utils/tests/test_ansi.py b/evennia/utils/tests/test_ansi.py index 4ff9d468c6..d8b0cac5bf 100644 --- a/evennia/utils/tests/test_ansi.py +++ b/evennia/utils/tests/test_ansi.py @@ -8,7 +8,16 @@ Test of the ANSI parsing and ANSIStrings. from django.test import TestCase -from evennia.utils.ansi import ANSIString as AN +from evennia.utils.ansi import ( + ANSIString as AN, + ANSI_RED, + ANSI_CYAN, + ANSI_YELLOW, + ANSI_GREEN, + ANSI_BLUE, + ANSI_HILITE, + ANSI_NORMAL, +) class TestANSIString(TestCase): @@ -20,7 +29,9 @@ class TestANSIString(TestCase): self.example_raw = "|relectric |cboogaloo|n" self.example_ansi = AN(self.example_raw) self.example_str = "electric boogaloo" - self.example_output = "\x1b[1m\x1b[31melectric \x1b[1m\x1b[36mboogaloo\x1b[0m" + self.example_output = ( + f"{ANSI_HILITE}{ANSI_RED}electric {ANSI_HILITE}{ANSI_CYAN}boogaloo{ANSI_NORMAL}" + ) def test_length(self): self.assertEqual(len(self.example_ansi), 17) @@ -52,3 +63,135 @@ class TestANSIString(TestCase): self.assertEqual(split2, split3, "Split 2 and 3 differ") self.assertEqual(split1, split2, "Split 1 and 2 differ") self.assertEqual(split1, split3, "Split 1 and 3 differ") + + def test_getitem_index_access(self): + """Test individual character access via indexing""" + # Test accessing individual characters + self.assertEqual(self.example_ansi[0].clean(), "e") + self.assertEqual(self.example_ansi[9].clean(), "b") + self.assertEqual(self.example_ansi[-1].clean(), "o") + self.assertEqual(self.example_ansi[-2].clean(), "o") + + # Verify ANSI codes are preserved when accessing characters + first_char = self.example_ansi[0] + self.assertTrue(isinstance(first_char, AN)) + # First character should have red color code + self.assertIn(ANSI_RED, first_char.raw()) + + # Test character at color boundary (first character after color change) + ninth_char = self.example_ansi[9] + self.assertEqual(ninth_char.clean(), "b") + # Should have cyan color code + self.assertIn(ANSI_CYAN, ninth_char.raw()) + + def test_getitem_slice_access(self): + """Test slice access""" + # Test basic slicing + substring = self.example_ansi[0:8] + self.assertEqual(substring.clean(), "electric") + self.assertTrue(isinstance(substring, AN)) + + # Test slicing with step + substring2 = self.example_ansi[9:17] + self.assertEqual(substring2.clean(), "boogaloo") + + # Test negative indices + last_three = self.example_ansi[-3:] + self.assertEqual(last_three.clean(), "loo") + + # Verify ANSI codes are preserved in slices + first_word = self.example_ansi[0:8] + self.assertIn(ANSI_RED, first_word.raw()) + + def test_getitem_edge_cases(self): + """Test edge cases for indexing""" + # Test with string with no ANSI codes + plain = AN("plain text") + self.assertEqual(plain[0].clean(), "p") + self.assertEqual(plain[6].clean(), "t") + + # Test with single character + single = AN("|rX|n") + self.assertEqual(len(single), 1) + self.assertEqual(single[0].clean(), "X") + + # Test IndexError + with self.assertRaises(IndexError): + _ = self.example_ansi[100] + + def test_upper_method(self): + """Test upper() method""" + # Test basic upper with ANSI codes + result = self.example_ansi.upper() + self.assertEqual(result.clean(), "ELECTRIC BOOGALOO") + self.assertTrue(isinstance(result, AN)) + + # Verify ANSI codes are preserved + self.assertIn(ANSI_RED, result.raw()) + self.assertIn(ANSI_CYAN, result.raw()) + + # Test with mixed case + mixed = AN("|rHeLLo |cWoRLd|n") + self.assertEqual(mixed.upper().clean(), "HELLO WORLD") + + def test_lower_method(self): + """Test lower() method""" + # Test basic lower with ANSI codes + upper_ansi = AN("|rELECTRIC |cBOOGALOO|n") + result = upper_ansi.lower() + self.assertEqual(result.clean(), "electric boogaloo") + self.assertTrue(isinstance(result, AN)) + + # Verify ANSI codes are preserved + self.assertIn(ANSI_RED, result.raw()) + self.assertIn(ANSI_CYAN, result.raw()) + + def test_capitalize_method(self): + """Test capitalize() method""" + # Test basic capitalize with ANSI codes + lower_ansi = AN("|relectric |cboogaloo|n") + result = lower_ansi.capitalize() + self.assertEqual(result.clean(), "Electric boogaloo") + self.assertTrue(isinstance(result, AN)) + + # Verify ANSI codes are preserved + self.assertIn(ANSI_RED, result.raw()) + + def test_swapcase_method(self): + """Test swapcase() method""" + # Test basic swapcase with ANSI codes + mixed = AN("|rElEcTrIc |cBoOgAlOo|n") + result = mixed.swapcase() + self.assertEqual(result.clean(), "eLeCtRiC bOoGaLoO") + self.assertTrue(isinstance(result, AN)) + + # Verify ANSI codes are preserved + self.assertIn(ANSI_RED, result.raw()) + self.assertIn(ANSI_CYAN, result.raw()) + + def test_transform_with_dense_ansi(self): + """Test string transformation with ANSI codes between every character""" + # Simulate rainbow text with ANSI between each character + dense = AN("|rh|ce|yl|gl|bo|n") + self.assertEqual(dense.clean(), "hello") + + # Test upper preserves all ANSI codes + upper_dense = dense.upper() + self.assertEqual(upper_dense.clean(), "HELLO") + self.assertTrue(isinstance(upper_dense, AN)) + + # Verify all color codes are still present + raw = upper_dense.raw() + self.assertIn(ANSI_RED, raw) + self.assertIn(ANSI_CYAN, raw) + self.assertIn(ANSI_YELLOW, raw) + self.assertIn(ANSI_GREEN, raw) + self.assertIn(ANSI_BLUE, raw) + + def test_transform_without_ansi(self): + """Test string transformation on plain strings""" + plain = AN("hello world") + + self.assertEqual(plain.upper().clean(), "HELLO WORLD") + self.assertEqual(plain.lower().clean(), "hello world") + self.assertEqual(plain.capitalize().clean(), "Hello world")