From b69ceb18d2739379b69c52459147341501e45694 Mon Sep 17 00:00:00 2001 From: jczetta Date: Sat, 14 Jul 2012 16:26:30 -0400 Subject: [PATCH] redefined __getattr__ so all str methods work All str methods work for ColoredStrings with expected functionality, excepting .join(). Retained some other method definitions in order to retain expected functionality. Added tests for ColoredString. --- clint/textui/colored.py | 14 +++++++++++--- test_clint.py | 26 ++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) 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()