Skip to content

Commit 0d7d533

Browse files
committed
Added __str__() back to ColorBase for more flexibility in its usage
Updated documentation and added unit tests
1 parent a73d300 commit 0d7d533

File tree

2 files changed

+36
-10
lines changed

2 files changed

+36
-10
lines changed

cmd2/ansi.py

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,27 @@ class ColorBase(Enum):
3232
Base class for fg and bg classes
3333
This expects the child classes to define enums of: color name -> ANSI color sequence
3434
"""
35+
def __str__(self) -> str:
36+
"""
37+
Return ANSI color sequence instead of enum name
38+
This is helpful when using a ColorBase in an f-string or format() call
39+
e.g. my_str = "{}hello{}".format(fg.blue, fg.reset)
40+
"""
41+
return self.value
42+
3543
def __add__(self, other: Any) -> str:
36-
"""Return self + other as string"""
37-
return self.value + other
44+
"""
45+
Support building a color string when self is the left operand
46+
e.g. fg.blue + "hello"
47+
"""
48+
return str(self) + other
3849

3950
def __radd__(self, other: Any) -> str:
40-
"""Return other + self as string"""
41-
return other + self.value
51+
"""
52+
Support building a color string when self is the right operand
53+
e.g. "hello" + fg.reset
54+
"""
55+
return other + str(self)
4256

4357
@classmethod
4458
def colors(cls) -> List[str]:

tests/test_ansi.py

Lines changed: 18 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -121,14 +121,26 @@ def test_async_alert_str(cols, prompt, line, cursor, msg, expected):
121121
assert alert_str == expected
122122

123123

124-
def test_fg_enum():
125-
assert ansi.fg_lookup('bright_red') == ansi.fg_lookup(ansi.fg.bright_red)
124+
def test_cast_color_as_str():
125+
assert str(ansi.fg.blue) == ansi.fg.blue.value
126+
assert str(ansi.bg.blue) == ansi.bg.blue.value
127+
128+
129+
def test_color_str_building():
130+
from cmd2.ansi import fg, bg
131+
assert fg.blue + "hello" == fg.blue.value + "hello"
132+
assert bg.blue + "hello" == bg.blue.value + "hello"
133+
assert fg.blue + "hello" + fg.reset == fg.blue.value + "hello" + fg.reset.value
134+
assert bg.blue + "hello" + bg.reset == bg.blue.value + "hello" + bg.reset.value
135+
assert fg.blue + bg.white + "hello" + fg.reset + bg.reset == \
136+
fg.blue.value + bg.white.value + "hello" + fg.reset.value + bg.reset.value
137+
126138

127-
def test_bg_enum():
139+
def test_color_enum():
140+
assert ansi.fg_lookup('bright_red') == ansi.fg_lookup(ansi.fg.bright_red)
128141
assert ansi.bg_lookup('green') == ansi.bg_lookup(ansi.bg.green)
129142

130-
def test_fg_colors():
131-
assert list(ansi.fg.__members__.keys()) == ansi.fg.colors()
132143

133-
def test_bg_colors():
144+
def test_colors_list():
145+
assert list(ansi.fg.__members__.keys()) == ansi.fg.colors()
134146
assert list(ansi.bg.__members__.keys()) == ansi.bg.colors()

0 commit comments

Comments
 (0)