Added tests for ANSIString, and fixed some bugs for it along the way.

src/settings_default.py

@@ -515,6 +515,7 @@ INSTALLED_APPS = (
     'src.comms',
     'src.help',
     'src.scripts',
+    'src.utils',
     'src.web.news',
     'src.web.website',)
 # The user profile extends the User object with more functionality;

src/utils/ansi.py

@@ -166,6 +166,12 @@ class ANSIParser(object):
                 else:
                     return ANSI_NORMAL + ANSI_BLUE
 
+    def strip_ansi(self, string):
+        """
+        Strips raw ANSI codes from a string.
+        """
+        return self.ansi_regex.sub("", string)
+
     def parse_ansi(self, string, strip_ansi=False, xterm256=False):
         """
         Parses a string, subbing color codes according to
@@ -199,7 +205,7 @@ class ANSIParser(object):
         if strip_ansi:
             # remove all ansi codes (including those manually
             # inserted in string)
-            parsed_string = self.ansi_regex.sub("", parsed_string)
+            return self.strip_ansi(string)
 
         # cache and crop old cache
         _PARSE_CACHE[cachekey] = parsed_string
@@ -378,7 +384,7 @@ def _on_raw(func_name):
                 args.insert(0, string)
         except IndexError:
             pass
-        result = _query_super(func_name)(self, *args, **kwargs)
+        result = getattr(self._raw_string, func_name)(*args, **kwargs)
         if isinstance(result, basestring):
             return ANSIString(result, decoded=True)
         return result
@@ -416,7 +422,7 @@ class ANSIMeta(type):
                 'rfind', 'rindex', '__len__']:
             setattr(cls, func_name, _query_super(func_name))
         for func_name in [
-                '__mul__', '__mod__', 'expandtabs', '__rmul__', 'join',
+                '__mul__', '__mod__', 'expandtabs', '__rmul__',
                 'decode', 'replace', 'format', 'encode']:
             setattr(cls, func_name, _on_raw(func_name))
         for func_name in [
@@ -456,13 +462,17 @@ class ANSIString(unicode):
         parser = kwargs.get('parser', ANSI_PARSER)
         decoded = kwargs.get('decoded', False) or hasattr(string, '_raw_string')
         if not decoded:
+            # Completely new ANSI String
+            clean_string = unicode(parser.parse_ansi(string, strip_ansi=True))
             string = parser.parse_ansi(string)
-        if hasattr(string, '_clean_string'):
+        elif hasattr(string, '_clean_string'):
+            # It's already an ANSIString
             clean_string = string._clean_string
             string = string._raw_string
         else:
-            clean_string = unicode(parser.parse_ansi(
+            # It's a string that has been pre-ansi decoded.
-                string, strip_ansi=True))
+            clean_string = parser.strip_ansi(string)
+
         if not isinstance(string, unicode):
             string = string.decode('utf-8')
         else:
@@ -491,7 +501,7 @@ class ANSIString(unicode):
         """
         return "ANSIString(%s, decoded=True)" % repr(self._raw_string)
 
-    def __init__(self, text="", parser=ANSI_PARSER, **kwargs):
+    def __init__(self, *_, **kwargs):
         """
         When the ANSIString is first initialized, a few internal variables
         have to be set.
@@ -517,8 +527,8 @@ class ANSIString(unicode):
         tables for which characters in the raw string are related to ANSI
         escapes, and which are for the readable text.
         """
-        self.parser = parser
+        self.parser = kwargs.pop('parser', ANSI_PARSER)
-        super(ANSIString, self).__init__(text)
+        super(ANSIString, self).__init__()
         self._code_indexes, self._char_indexes = self._get_indexes()
 
     def __add__(self, other):
@@ -583,8 +593,8 @@ class ANSIString(unicode):
                 string += self._raw_string[i]
             except IndexError:
                 pass
-        if slc.stop is not None:
+        if i is not None:
-            append_tail = self._get_interleving(slc.stop)
+            append_tail = self._get_interleving(self._char_indexes.index(i) + 1)
         else:
             append_tail = ''
         return ANSIString(string + append_tail, decoded=True)
@@ -721,11 +731,11 @@ class ANSIString(unicode):
             if next < 0:
                 break
             # Get character codes after the index as well.
-            res.append(self[start:next] + self._get_interleving(next))
+            res.append(self[start:next])
             start = next + bylen
             maxsplit -= 1   # NB. if it's already < 0, it stays < 0
 
-        res.append(self[start:len(self)] + self._get_interleving(len(self)))
+        res.append(self[start:len(self)])
         return res
 
     def rsplit(self, by, maxsplit=-1):
@@ -747,14 +757,28 @@ class ANSIString(unicode):
             if next < 0:
                 break
             # Get character codes after the index as well.
-            res.append(self[next+bylen:end] + self._get_interleving(end))
+            res.append(self[next+bylen:end])
             end = next
             maxsplit -= 1   # NB. if it's already < 0, it stays < 0
 
-        res.append(self[:end] + self._get_interleving(end))
+        res.append(self[:end])
         res.reverse()
         return res
 
+    def join(self, iterable):
+        """
+        Joins together strings in an iterable.
+        """
+        result = ANSIString('')
+        last_item = None
+        for item in iterable:
+            if last_item is not None:
+                result += self
+            result += item
+            last_item = item
+        return result
+
+
     @_spacing_preflight
     def center(self, width, fillchar, difference):
         """

src/utils/models.py

@@ -0,0 +1,4 @@
+"""
+Dummy models.py file to allow us to add utils to the apps list so Django's
+test runner will recognize it.
+"""

src/utils/tests.py

@@ -0,0 +1,112 @@
+import re
+
+try:
+    from django.utils.unittest import TestCase
+except ImportError:
+    from django.test import TestCase
+
+from ansi import ANSIString
+
+
+class ANSIStringTestCase(TestCase):
+    def checker(self, ansi, raw, clean):
+        """
+        Verifies the raw and clean strings of an ANSIString match expected
+        output.
+        """
+        self.assertEqual(unicode(ansi.clean()), clean)
+        self.assertEqual(unicode(ansi.raw()), raw)
+
+    def table_check(self, ansi, char, code):
+        """
+        Verifies the indexes in an ANSIString match what they should.
+        """
+        self.assertEqual(ansi._char_indexes, char)
+        self.assertEqual(ansi._code_indexes, code)
+
+    def test_instance(self):
+        """
+        Make sure the ANSIString is always constructed correctly.
+        """
+        clean = u'This isA{r testTest'
+        encoded = u'\x1b[1m\x1b[32mThis is\x1b[1m\x1b[31mA{r test\x1b[0mTest\x1b[0m'
+        target = ANSIString(r'{gThis is{rA{{r test{nTest{n')
+        char_table = [9, 10, 11, 12, 13, 14, 15, 25, 26, 27, 28, 29, 30, 31,
+                      32, 37, 38, 39, 40]
+        code_table = [0, 1, 2, 3, 4, 5, 6, 7, 8, 16, 17, 18, 19, 20, 21, 22,
+                      23, 24, 33, 34, 35, 36, 41, 42, 43, 44]
+        self.checker(target, encoded, clean)
+        self.table_check(target, char_table, code_table)
+        self.checker(ANSIString(target), encoded, clean)
+        self.table_check(ANSIString(target), char_table, code_table)
+        self.checker(ANSIString(encoded, decoded=True), encoded, clean)
+        self.table_check(ANSIString(encoded, decoded=True), char_table,
+                         code_table)
+        self.checker(ANSIString('Test'), u'Test', u'Test')
+        self.table_check(ANSIString('Test'), [0, 1, 2, 3], [])
+        self.checker(ANSIString(''), u'', u'')
+
+    def test_slice(self):
+        """
+        Verifies that slicing an ANSIString results in expected color code
+        distribution.
+        """
+        target = ANSIString(r'{gTest{rTest{n')
+        result = target[:3]
+        self.checker(result, u'\x1b[1m\x1b[32mTes', u'Tes')
+        result = target[:4]
+        self.checker(result, u'\x1b[1m\x1b[32mTest\x1b[1m\x1b[31m', u'Test')
+        result = target[:]
+        self.checker(
+            result,
+            u'\x1b[1m\x1b[32mTest\x1b[1m\x1b[31mTest\x1b[0m',
+            u'TestTest')
+        result = target[:-1]
+        self.checker(
+            result,
+            u'\x1b[1m\x1b[32mTest\x1b[1m\x1b[31mTes',
+            u'TestTes')
+        result = target[0:0]
+        self.checker(
+            result,
+            u'',
+            u'')
+
+    def test_split(self):
+        """
+        Verifies that re.split and .split behave similarly and that color
+        codes end up where they should.
+        """
+        target = ANSIString("{gThis is {nA split string{g")
+        first = (u'\x1b[1m\x1b[32mThis is \x1b[0m', u'This is ')
+        second = (u'\x1b[1m\x1b[32m\x1b[0m split string\x1b[1m\x1b[32m',
+                  u' split string')
+        re_split = re.split('A', target)
+        normal_split = target.split('A')
+        self.assertEqual(re_split, normal_split)
+        self.assertEqual(len(normal_split), 2)
+        self.checker(normal_split[0], *first)
+        self.checker(normal_split[1], *second)
+
+    def test_join(self):
+        """
+        Verify that joining a set of ANSIStrings works.
+        """
+        # This isn't the desired behavior, but the expected one. Python
+        # concatinates the in-memory representation with the built-in string's
+        # join.
+        l = [ANSIString("{gTest{r") for s in range(0, 3)]
+        # Force the generator to be evaluated.
+        result = "".join(l)
+        self.assertEqual(unicode(result), u'TestTestTest')
+        result = ANSIString("").join(l)
+        self.checker(result, u'\x1b[1m\x1b[32mTest\x1b[1m\x1b[31m\x1b[1m\x1b'
+                             u'[32mTest\x1b[1m\x1b[31m\x1b[1m\x1b[32mTest'
+                             u'\x1b[1m\x1b[31m', u'TestTestTest')
+
+    def test_len(self):
+        """
+        Make sure that length reporting on ANSIStrings does not include
+        ANSI codes.
+        """
+        self.assertEqual(len(ANSIString('{gTest{n')), 4)