Optimize ANSIString

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]
+        char_pos = self._char_indexes[item]
-        result = ""
+        clean = self._raw_string[char_pos]
-        # Get the character they're after, and replay all escape sequences
+
-        # previous to it.
+        code_indexes_set = set(self._code_indexes)
-        for index in range(0, item + 1):
+
-            if index in self._code_indexes:
+        result_chars = [
-                result += self._raw_string[index]
+            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

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 (tests optimized __getitem__)"""
+        # 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 (tests optimized __getitem__ via _slice)"""
+        # 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 (uses optimized _transform)"""
+        # 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 (uses optimized _transform)"""
+        # 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 (uses optimized _transform)"""
+        # 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 (uses optimized _transform)"""
+        # 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")