Fix errors with unknown tags in parse() (#4) and strip() (#6)

ansimarkup/__init__.py

@@ -1,9 +1,9 @@
-from . markup import AnsiMarkup, AnsiMarkupError, MismatchedTag
+from . markup import AnsiMarkup, AnsiMarkupError, MismatchedTag, UnbalancedTag
 
 
 _ansimarkup = AnsiMarkup()
 parse = _ansimarkup.parse
 ansiprint = _ansimarkup.ansiprint
 
 
-__all__ = 'AnsiMarkup', 'AnsiMarkupError', 'MismatchedTag', 'parse', 'ansiprint'
+__all__ = 'AnsiMarkup', 'AnsiMarkupError', 'MismatchedTag', 'UnbalancedTag', 'parse', 'ansiprint'

ansimarkup/markup.py

@@ -14,6 +14,9 @@ class AnsiMarkupError(Exception):
 class MismatchedTag(AnsiMarkupError):
     pass
 
+class UnbalancedTag(AnsiMarkupError):
+    pass
+
 
 class AnsiMarkup:
     '''
@@ -35,14 +38,13 @@ def __init__(self, tags=None, always_reset=False, tag_sep='<>'):
         self.user_tags = tags if tags else {}
         self.always_reset = always_reset
 
-        self.re_tag_start, self.re_tag_end = self.compile_tag_regex(tag_sep)
+        self.re_tag = self.compile_tag_regex(tag_sep)
 
     def parse(self, text):
         '''Return a string with markup tags converted to ansi-escape sequences.'''
         tags, results = [], []
 
-        text = self.re_tag_start.sub(lambda m: self.sub_start(m, tags, results), text)
-        text = self.re_tag_end.sub(lambda m: self.sub_end(m, tags, results), text)
+        text = self.re_tag.sub(lambda m: self.sub_tag(m, tags, results), text)
 
         if self.always_reset:
             if not text.endswith(Style.RESET_ALL):
@@ -57,16 +59,28 @@ def ansiprint(self, *args, **kwargs):
 
     def strip(self, text):
         '''Return string with markup tags removed.'''
-        # TODO: This also strips unknown tags.
-        text = self.re_tag_start.sub('', text)
-        text = self.re_tag_end.sub('', text)
+        tags, results = [], []
+
+        text = self.re_tag.sub(lambda m: self.clear_tag(m, tags, results), text)
+
         return text
 
     def __call__(self, text):
         return self.parse(text)
 
-    def sub_start(self, match, tag_list, res_list):
-        tag = match.group(1)
+    def sub_tag(self, match, tag_list, res_list):
+        markup, tag = match.group(0), match.group(1)
+        closing = markup[1] == '/'
+        res = None
+
+        # Early exit if the closing tag matches the last known opening tag
+        if closing and tag_list and tag_list[-1] == tag:
+            tag_list.pop()
+            res_list.pop()
+
+            res = Style.RESET_ALL + ''.join(res_list)
+
+            return res
 
         # User-defined tags take preference over all other.
         if tag in self.user_tags:
@@ -77,85 +91,90 @@ def sub_start(self, match, tag_list, res_list):
         elif tag in all_tags:
             res = all_tags[tag]
 
-        # An alternative syntax for setting the foreground color (e.g. <fg red>).
-        elif tag.startswith('fg '):
-            color = tag[3:]
-            if color.isdigit() and int(color) <= 255:
-                res = '\033[38;5;%sm' % color
-            elif color.startswith('#'):
-                res = '\033[38;2;%s;%s;%sm' % hex_to_rgb(color[1:])
+        # An alternative syntax for setting the color (e.g. <fg red>, <bg red>).
+        elif tag.startswith('fg ') or tag.startswith('bg '):
+            st, color = tag[:2], tag[3:]
+            code = '38' if st == 'fg' else '48'
+
+            if st == 'fg' and color in foreground:
+                res = foreground[color]
+            elif st == 'bg' and color.islower() and color.upper() in background:
+                res = background[color.upper()]
+            elif color.isdigit() and int(color) <= 255:
+                res = '\033[%s;5;%sm' % (code, color)
+            elif re.match(r'#(?:[a-fA-F0-9]{3}){1,2}$', color):
+                hex_color = color[1:]
+                if len(hex_color) == 3:
+                    hex_color *= 2
+                res = '\033[%s;2;%s;%s;%sm' % ((code,) + hex_to_rgb(hex_color))
             elif color.count(',') == 2:
-                res = '\033[38;2;%s;%s;%sm' % tuple(color.split(','))
-            else:
-                res = foreground.get(color, match.group())
-
-        # An alternative syntax for setting the background color (e.g. <bg red>).
-        elif tag.startswith('bg '):
-            color = tag[3:]
-            if color.isdigit() and int(color) <= 255:
-                res = '\033[48;5;%sm' % color
-            elif color.startswith('#'):
-                res = '\033[48;2;%s;%s;%sm' % hex_to_rgb(color[1:])
-            elif color.count(',') == 2:
-                res = '\033[48;2;%s;%s;%sm' % tuple(color.split(','))
-            else:
-                res = background.get(color.upper(), match.group())
+                colors = tuple(color.split(','))
+                if all(x.isdigit() and int(x) <= 255 for x in colors):
+                    res = '\033[%s;2;%s;%s;%sm' % ((code,) + colors)
 
         # Shorthand formats (e.g. <red,blue>, <bold,red,blue>).
         elif ',' in tag:
             el_count = tag.count(',')
 
             if el_count == 1:
                 fg, bg = tag.split(',')
-                res = foreground.get(fg, '') + background.get(bg.upper(), '')
+                if fg in foreground and bg.islower() and bg.upper() in background:
+                    res = foreground[fg] + background[bg.upper()]
 
             elif el_count == 2:
                 st, fg, bg = tag.split(',')
-                res = style[st] + foreground.get(fg, '') + background.get(bg.upper(), '')
+                if st in style and (fg != '' or bg != ''):
+                    if fg == '' or fg in foreground:
+                        if bg == '' or (bg.islower() and bg.upper() in background):
+                            res = style[st] + foreground.get(fg, '') + background.get(bg.upper(), '')
 
         # If nothing matches, return the full tag (i.e. <unknown>text</...>).
-        else:
-            return match.group()
+        if res is None:
+            return markup
+
+        # If closing tag is known, but did not early exit, something is wrong
+        if closing:
+            if tag in tag_list:
+                raise UnbalancedTag('closing tag "%s" violates nesting rules.' % markup)
+            else:
+                raise MismatchedTag('closing tag "%s" has no corresponding opening tag' % markup)
 
         tag_list.append(tag)
         res_list.append(res)
-        return res
 
-    def sub_end(self, match, tag_list, res_list):
-        tag = match.group(1)
+        return res
 
-        if tag_list:
-            try:
-                idx = tag_list.index(tag)
-            except ValueError:
-                raise MismatchedTag('closing tag "</%s>" has no corresponding opening tag' % tag)
+    def clear_tag(self, match, tag_list, res_list):
+        pre_length = len(tag_list)
+        try:
+            self.sub_tag(match, tag_list, res_list)
 
-            res = ''.join(res_list[:idx])
+            # If list did not change, the tag is unknown
+            if len(tag_list) == pre_length:
+                return match.group(0)
 
-            del tag_list[idx]
-            del res_list[idx]
-            return Style.RESET_ALL + res
-        else:
-            return Style.RESET_ALL
+            # Otherwise, tag matched so remove it
+            else:
+                return ''
 
-        return match.group()
+        # Tag matched but is invalid, remove it anyway
+        except (UnbalancedTag, MismatchedTag):
+            return ''
 
     def compile_tag_regex(self, tag_sep):
         # Optimize the default:
         if tag_sep == '<>':
-            tag_start = re.compile(r'<([^/>]+)>')
-            tag_end   = re.compile(r'</([^>]+)>')
-            return tag_start, tag_end
+            tag_regex = re.compile(r'</?([^>]+)>')
+            return tag_regex
 
         if len(tag_sep) < 2:
             raise ValueError('tag_sep needs to have at least two element (e.g. "<>")')
 
         if tag_sep[0] == tag_sep[1]:
             raise ValueError('opening and closing characters cannot be the same')
 
-        tag_start = r'{0}([^/{1}]+){1}'.format(tag_sep[0], tag_sep[1])
-        tag_end   = r'{0}/([^{1}]+){1}'.format(tag_sep[0], tag_sep[1])
-        return re.compile(tag_start), re.compile(tag_end)
+        tag_regex = r'{0}/?([^{1}]+){1}'.format(tag_sep[0], tag_sep[1])
+        return re.compile(tag_regex)
 
 
 def hex_to_rgb(value):

tests/test_markup.py

@@ -23,36 +23,134 @@ def test_flat_colors():
 	assert p('<b,r,y>1</b,r,y>') == p('<bold,red,yellow>1</bold,red,yellow>') == S.BRIGHT + F.RED + B.YELLOW + '1' + S.RESET_ALL
 	assert p('<b,r,>1</b,r,>') == p('<bold,red,>1</bold,red,>') == S.BRIGHT + F.RED + '1' + S.RESET_ALL
 
+	assert p('<fg RED>1</fg RED>') == '<fg RED>1</fg RED>'
+	assert p('<fg light-blue2>1</fg light-blue2>') == '<fg light-blue2>1</fg light-blue2>'
+
+	assert p('<bg ,red>1</bg ,red>') == '<bg ,red>1</bg ,red>'
+	assert p('<bg red,>1</bg red,>') == '<bg red,>1</bg red,>'
+	assert p('<bg a,z>1</bg a,z>') == '<bg a,z>1</bg a,z>'
+	assert p('<bg blue,yelllow>1</bg blue,yelllow>') == '<bg blue,yelllow>1</bg blue,yelllow>'
+	assert p('<bg r, y>1</bg r, y>') == '<bg r, y>1</bg r, y>'
+
+	assert p('<>1</>') == '<>1</>'
+	assert p('</>1</>') == '</>1</>'
+	assert p('<,>1</,>') == '<,>1</,>'
+	assert p('<,,>1</,,>') == '<,,>1</,,>'
+	assert p('<z,z>1</z,z>') == '<z,z>1</z,z>'
+	assert p('<z,z,z>1</z,z,z>') == '<z,z,z>1</z,z,z>'
+
+	assert p('<red,RED>1</red,RED>') == '<red,RED>1</red,RED>'
+	assert p('<bold,red>1</bold,red>') == '<bold,red>1</bold,red>'
+	assert p('<b,red,RED>1</b,red,RED>') == '<b,red,RED>1</b,red,RED>'
+	assert p('<b,red,bold>1</b,red,bold>') == '<b,red,bold>1</b,red,bold>'
+	assert p('<b,,>1</b,,>') == '<b,,>1</b,,>'
+	assert p('<b,a,>1</b,a,>') == '<b,a,>1</b,a,>'
+	assert p('<b,,z>1</b,,z>') == '<b,,z>1</b,,z>'
+
 
 def test_xterm_color():
 	assert p('<fg 200>1</fg 200>') == '\x1b[38;5;200m' '1' '\x1b[0m'
 	assert p('<bg 100><fg 200>1')  == '\x1b[48;5;100m' '\x1b[38;5;200m' '1'
+
 	assert p('<fg 256>1')          == '<fg 256>1'
+	assert p('<bg 256>1</bg 256>') == '<bg 256>1</bg 256>'
+	assert p('<bg 12 >1</bg 12 >') == '<bg 12 >1</bg 12 >'
+	assert p('<bg 2222>1</bg 2222>') == '<bg 2222>1</bg 2222>'
 
 
 def test_xterm_hex():
-	assert p('<fg #ff0000>1') == '\x1b[38;2;255;0;0m' '1'
+	assert p('<fg #ff0000>1') == p('<fg #FF0000>1') == '\x1b[38;2;255;0;0m' '1'
 	assert p('<bg #00A000><fg #ff0000>1') == '\x1b[48;2;0;160;0m' '\x1b[38;2;255;0;0m' '1'
+	assert p('<fg #F12>1</fg #F12>') == p('<fg #F12F12>1</fg #F12F12>') == '\x1b[38;2;241;47;18m' + '1' + S.RESET_ALL
+
+	assert p('<fg #>1</fg #>') == '<fg #>1</fg #>'
+	assert p('<bg #12>1</bg #12>') == '<bg #12>1</bg #12>'
+	assert p('<fg #1234567>1</fg #1234567>') == '<fg #1234567>1</fg #1234567>'
+	assert p('<bg #E7G>1</bg #E7G>') == '<bg #E7G>1</bg #E7G>'
+	assert p('<fg #F2D1GZ>1</fg #F2D1GZ>') == '<fg #F2D1GZ>1</fg #F2D1GZ>'
 
 
 def test_xterm_rgb():
 	assert p('<fg 255,0,0>1') == p('<fg #ff0000>1') == '\x1b[38;2;255;0;0m' '1'
 	assert p('<bg 0,160,0><fg 255,0,0>1') == p('<bg #00A000><fg #ff0000>1') == '\x1b[48;2;0;160;0m' '\x1b[38;2;255;0;0m' '1'
 
+	assert p('<fg>1</fg>') == '<fg>1</fg>'
+	assert p('<fg >1</fg >') == '<fg >1</fg >'
+	assert p('<fg ,>1</fg ,>') == '<fg ,>1</fg ,>'
+	assert p('<fg ,,>1</fg ,,>') == '<fg ,,>1</fg ,,>'
+	assert p('<fg ,,,>1</fg ,,,>') == '<fg ,,,>1</fg ,,,>'
+	assert p('<fg 1,2,>1</fg 1,2,>') == '<fg 1,2,>1</fg 1,2,>'
+	assert p('<fg 256,120,120>1</fg 256,120,120>') == '<fg 256,120,120>1</fg 256,120,120>'
+
 
 def test_nested():
 	assert p('0<b>1<d>2</d>3</b>4') == '0' + S.BRIGHT + '1' + S.DIM + '2' + S.RESET_ALL + S.BRIGHT + '3' + S.RESET_ALL + '4'
+	assert p('<d>0<b>1<d>2</d>3</b>4</d>') == S.DIM + '0' + S.BRIGHT + '1' + S.DIM + '2' + S.RESET_ALL + S.DIM + S.BRIGHT + '3' + S.RESET_ALL + S.DIM +'4' + S.RESET_ALL
 
 
-def test_errors():
+def test_mismatched_error():
 	with raises(MismatchedTag):
 		p('<b>1</d>')
 
+	with raises(MismatchedTag):
+		p('</b>')
+
+	with raises(MismatchedTag):
+		p('<b>1</b></b>')
+
+	with raises(MismatchedTag):
+		p('<red><b>1</b></b></red>')
+
+	with raises(MismatchedTag):
+		p('<tag>1</b>')
+
+
+def test_unbalanced_error():
+	with raises(UnbalancedTag):
+		p('<r><Y>1</r>2</Y>')
+
+	with raises(UnbalancedTag):
+		p('<r><r><Y>1</r>2</Y></r>')
+
+	with raises(UnbalancedTag):
+		p('<r><Y><r></r></r></Y>')
+
+
+def test_unknown_tags():
+	assert p('<tag>1</tag>') == '<tag>1</tag>'
+	assert p('<tag>') == '<tag>'
+	assert p('</tag>') == '</tag>'
+
+	assert p('<b>1</b><tag>2</tag>') == S.BRIGHT + '1' + S.RESET_ALL + '<tag>2</tag>'
+	assert p('<tag>1</tag><b>2</b>') == '<tag>1</tag>' + S.BRIGHT + '2' + S.RESET_ALL
+
+	assert p('<b>1</b><tag>2</tag><b>3</b>') == S.BRIGHT + '1' + S.RESET_ALL + '<tag>2</tag>' + S.BRIGHT + '3' + S.RESET_ALL
+	assert p('<tag>1</tag><b>2</b><tag>3</tag>') == '<tag>1</tag>' + S.BRIGHT + '2' + S.RESET_ALL + '<tag>3</tag>'
+
+	assert p('<b><tag>1</tag></b>') == S.BRIGHT + '<tag>1</tag>' + S.RESET_ALL
+	assert p('<tag><b>1</b></tag>') == '<tag>' + S.BRIGHT + '1' + S.RESET_ALL + '</tag>'
+
+	assert p('<b><tag>1</tag>') == S.BRIGHT + '<tag>1</tag>'
+	assert p('<tag>1</tag><b>') == '<tag>1</tag>' + S.BRIGHT
+
+	assert p('<b><tag>') == S.BRIGHT + '<tag>'
+	assert p('<tag><b>') == "<tag>" + S.BRIGHT
+
+	assert p('<b></tag>') == S.BRIGHT + '</tag>'
+	assert p('</tag><b>') == "</tag>" + S.BRIGHT
+
 
 def test_strip(am):
 	assert am.strip('<b>1</b>2<d>3</d>') == '123'
 	assert am.strip('<bold,red,yellow>1</bold,red,yellow>') == '1'
+	assert am.strip('<r><tag>1</tag></r>') == '<tag>1</tag>'
+	assert am.strip('<tag><r>1</r></tag>') == '<tag>1</tag>'
+	assert am.strip('<tag><r>1</tag></r>') == '<tag>1</tag>'
+	assert am.strip('<r><b>1</r></b>') == '1'
+	assert am.strip('<fg red>1<fg red>') == '1'
 
+	am = AnsiMarkup(tags={'red':'<red>', 'b,g,r':'<b,g,r>', 'fg 1,2,3': '</fg 1,2,3>'})
+	assert am.strip('<red>1</red><b,g,r>2</b,g,r><fg 1,2,3>3</fg 1,2,3>') == '123'
 
 def test_options():
 	am = AnsiMarkup(always_reset=True)
@@ -87,7 +185,3 @@ def test_tag_chars():
 	r1 = p('0<b>1<d>2</d>3</b>4')
 	r2 = am.parse('0{b}1{d}2{/d}3{/b}4')
 	assert r1 == r2 == '0' + S.BRIGHT + '1' + S.DIM + '2' + S.RESET_ALL + S.BRIGHT + '3' + S.RESET_ALL + '4'
-
-@mark.xfail
-def test_limitations():
-	assert p('<r><Y>1</r>2</Y>') == F.RED + B.YELLOW + '1' + S.RESET_ALL + B.YELLOW + '2' + S.RESET_ALL