Added __str__() back to ColorBase for more flexibility in its usage

Updated documentation and added unit tests

2 files changed, 36 insertions, 10 deletions

diff --git a/cmd2/ansi.py b/cmd2/ansi.py
index 27086d4a..3a106776 100644
--- a/cmd2/ansi.py
+++ b/cmd2/ansi.py
@@ -32,13 +32,27 @@ class ColorBase(Enum):
     Base class for fg and bg classes
     This expects the child classes to define enums of: color name -> ANSI color sequence
     """
+    def __str__(self) -> str:
+        """
+        Return ANSI color sequence instead of enum name
+        This is helpful when using a ColorBase in an f-string or format() call
+        e.g. my_str = "{}hello{}".format(fg.blue, fg.reset)
+        """
+        return self.value
+
     def __add__(self, other: Any) -> str:
-        """Return self + other as string"""
-        return self.value + other
+        """
+        Support building a color string when self is the left operand
+        e.g. fg.blue + "hello"
+        """
+        return str(self) + other
 
     def __radd__(self, other: Any) -> str:
-        """Return other + self as string"""
-        return other + self.value
+        """
+        Support building a color string when self is the right operand
+        e.g. "hello" + fg.reset
+        """
+        return other + str(self)
 
     @classmethod
     def colors(cls) -> List[str]:
diff --git a/tests/test_ansi.py b/tests/test_ansi.py
index 526a06e4..70ab5cb2 100644
--- a/tests/test_ansi.py
+++ b/tests/test_ansi.py
@@ -121,14 +121,26 @@ def test_async_alert_str(cols, prompt, line, cursor, msg, expected):
     assert alert_str == expected
 
 
-def test_fg_enum():
-    assert ansi.fg_lookup('bright_red') == ansi.fg_lookup(ansi.fg.bright_red)
+def test_cast_color_as_str():
+    assert str(ansi.fg.blue) == ansi.fg.blue.value
+    assert str(ansi.bg.blue) == ansi.bg.blue.value
+
+
+def test_color_str_building():
+    from cmd2.ansi import fg, bg
+    assert fg.blue + "hello" == fg.blue.value + "hello"
+    assert bg.blue + "hello" == bg.blue.value + "hello"
+    assert fg.blue + "hello" + fg.reset == fg.blue.value + "hello" + fg.reset.value
+    assert bg.blue + "hello" + bg.reset == bg.blue.value + "hello" + bg.reset.value
+    assert fg.blue + bg.white + "hello" + fg.reset + bg.reset == \
+           fg.blue.value + bg.white.value + "hello" + fg.reset.value + bg.reset.value
+
 
-def test_bg_enum():
+def test_color_enum():
+    assert ansi.fg_lookup('bright_red') == ansi.fg_lookup(ansi.fg.bright_red)
     assert ansi.bg_lookup('green') == ansi.bg_lookup(ansi.bg.green)
 
-def test_fg_colors():
-    assert list(ansi.fg.__members__.keys()) == ansi.fg.colors()
 
-def test_bg_colors():
+def test_colors_list():
+    assert list(ansi.fg.__members__.keys()) == ansi.fg.colors()
     assert list(ansi.bg.__members__.keys()) == ansi.bg.colors()