diff --git a/clint/textui/colored.py b/clint/textui/colored.py index 0906dde..2c18b75 100644 --- a/clint/textui/colored.py +++ b/clint/textui/colored.py @@ -45,6 +45,17 @@ class ColoredString(object): self.s = s self.color = color + def __getattr__(self, att): + def func_help(*args, **kwargs): + result = getattr(self.s, att)(*args, **kwargs) + if isinstance(result, basestring): + return self._new(result) + elif isinstance(result, list): + return [self._new(x) for x in result] + else: + return result + return func_help + @property def color_str(self): if sys.stdout.isatty() and not DISABLE_COLOR: @@ -84,9 +95,6 @@ class ColoredString(object): def __mul__(self, other): return (self.color_str * other) - def split(self, sep=None): - return [self._new(s) for s in self.s.split(sep)] - def _new(self, s): return ColoredString(self.color, s) diff --git a/test_clint.py b/test_clint.py index e9e0b9f..5bdeda2 100755 --- a/test_clint.py +++ b/test_clint.py @@ -16,5 +16,31 @@ class TablibTestCase(unittest.TestCase): def tearDown(self): pass +class ColoredStringTestCase(unittest.TestCase): + + def setUp(self): + from clint.textui.colored import ColoredString + + def tearDown(self): + pass + + def test_split(self): + from clint.textui.colored import ColoredString + new_str = ColoredString('red', "hello world") + output = new_str.split() + assert output[0].s == "hello" + + def test_find(self): + from clint.textui.colored import ColoredString + new_str = ColoredString('blue', "hello world") + output = new_str.find('h') + self.assertEqual(output, 0) + + def test_replace(self): + from clint.textui.colored import ColoredString + new_str = ColoredString('green', "hello world") + output = new_str.replace("world", "universe") + assert output.s == "hello universe" + if __name__ == '__main__': unittest.main()