colors: bug fix on new colorizing markup, and tests

Author: koreno

easypy/colors.py

@@ -135,9 +135,9 @@ def __new__(cls, text):
                 continue
             match = RE_PARSE_COLOR_MARKUP.match(part)
             if match:
-                stl, part1, part2 = match.groups()
+                stl, *parts = match.groups()
                 stl = stl.strip("_")
-                part = part1 or part2 or ""
+                part = next(filter(None, parts))
                 for l in part.splitlines():
                     self.tokens.append(self.ColoredToken(l, stl))
                     self.tokens.append(self.Token("\n"))
@@ -323,7 +323,7 @@ def uncolored(text, ansi=True, markup=True):
     if ansi:
         text = re.sub(re.escape(ANSI_BEGIN) + '.+?m', "", text)
     if markup:
-        text = RE_PARSE_COLOR_MARKUP.sub(lambda m: m.group(2) or m.group(3), text)
+        text = RE_PARSE_COLOR_MARKUP.sub(lambda m: next(filter(None, m.groups()[1:])), text)
     return text
 
 

tests/test_colors.py

@@ -0,0 +1,24 @@
+def test_colors():
+    from easypy.colors import Colorized, uncolored, colorize, register_colorizers
+
+    register_colorizers(bad=("red", "blue"))
+
+    opts = {
+        str(Colorized("RED(BLUE)<<XXX>>")),
+        str(Colorized("RED(BLUE)@[XXX]@")),
+        str(Colorized("RED(BLUE)@{XXX}@")),
+
+        str(colorize("RED(BLUE)<<XXX>>")),
+        str(colorize("RED(BLUE)@[XXX]@")),
+        str(colorize("RED(BLUE)@{XXX}@")),
+
+        str(colorize("BAD<<XXX>>")),
+        str(colorize("BAD@[XXX]@")),
+        str(colorize("BAD@{XXX}@")),
+    }
+
+    assert len(opts) == 1
+    [ret] = opts
+
+    assert ret == "\x1b[1;44;31mXXX\x1b[0m"
+    assert uncolored(ret) == "XXX"