Updated plumbum example

2 files changed, 43 insertions, 41 deletions

diff --git a/cmd2/ansi.py b/cmd2/ansi.py
index 3a106776..f54abf0a 100644
--- a/cmd2/ansi.py
+++ b/cmd2/ansi.py
@@ -3,9 +3,9 @@
 Support for ANSI escape sequences which are used for things like applying style to text,
 setting the window title, and asynchronous alerts.
  """
-from enum import Enum, unique
 import functools
 import re
+from enum import Enum
 from typing import Any, IO, List, Union
 
 import colorama
@@ -29,8 +29,10 @@ ANSI_STYLE_RE = re.compile(r'\x1b\[[^m]*m')
 
 class ColorBase(Enum):
     """
-    Base class for fg and bg classes
-    This expects the child classes to define enums of: color name -> ANSI color sequence
+    Base class used for defining color enums. See fg and bg classes for examples.
+    This expects the child classes to define enums of the follow structure
+        key: color name (e.g. black)
+        value: anything that when cast to a string returns an ANSI sequence
     """
     def __str__(self) -> str:
         """
@@ -38,7 +40,7 @@ class ColorBase(Enum):
         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
+        return str(self.value)
 
     def __add__(self, other: Any) -> str:
         """
@@ -57,12 +59,12 @@ class ColorBase(Enum):
     @classmethod
     def colors(cls) -> List[str]:
         """Return a list of color names."""
-        return [color.name for color in cls]
+        # Use __members__ to ensure we get all key names, including those which are aliased
+        return [color for color in cls.__members__]
 
 
 # Foreground colors
-# noinspection PyPep8Naming,DuplicatedCode
-@unique
+# noinspection PyPep8Naming
 class fg(ColorBase):
     """Enum class for foreground colors"""
     black = Fore.BLACK
@@ -85,8 +87,7 @@ class fg(ColorBase):
 
 
 # Background colors
-# noinspection PyPep8Naming,DuplicatedCode
-@unique
+# noinspection PyPep8Naming
 class bg(ColorBase):
     """Enum class for background colors"""
     black = Back.BLACK
diff --git a/examples/plumbum_colors.py b/examples/plumbum_colors.py
index 94815f50..2887ef1f 100755
--- a/examples/plumbum_colors.py
+++ b/examples/plumbum_colors.py
@@ -31,36 +31,37 @@ import cmd2
 from cmd2 import ansi
 from plumbum.colors import fg, bg
 
-FG_COLORS = {
-    'black': fg.Black,
-    'red': fg.DarkRedA,
-    'green': fg.MediumSpringGreen,
-    'yellow': fg.LightYellow,
-    'blue': fg.RoyalBlue1,
-    'magenta': fg.Purple,
-    'cyan': fg.SkyBlue1,
-    'white': fg.White,
-    'purple': fg.Purple,
-}
-
-BG_COLORS = {
-    'black': bg.BLACK,
-    'red': bg.DarkRedA,
-    'green': bg.MediumSpringGreen,
-    'yellow': bg.LightYellow,
-    'blue': bg.RoyalBlue1,
-    'magenta': bg.Purple,
-    'cyan': bg.SkyBlue1,
-    'white': bg.White,
-}
-
-
-def get_fg(fg):
-    return str(FG_COLORS[fg])
-
-
-def get_bg(bg):
-    return str(BG_COLORS[bg])
+
+class FgColors(ansi.ColorBase):
+    black = fg.Black
+    red = fg.DarkRedA
+    green = fg.MediumSpringGreen
+    yellow = fg.LightYellow
+    blue = fg.RoyalBlue1
+    magenta = fg.Purple
+    cyan = fg.SkyBlue1
+    white = fg.White
+    purple = fg.Purple
+
+
+class BgColors(ansi.ColorBase):
+    black = bg.BLACK
+    red = bg.DarkRedA
+    green = bg.MediumSpringGreen
+    yellow = bg.LightYellow
+    blue = bg.RoyalBlue1
+    magenta = bg.Purple
+    cyan = bg.SkyBlue1
+    white = bg.White
+    purple = bg.Purple
+
+
+def get_fg(name: str):
+    return str(FgColors[name])
+
+
+def get_bg(name: str):
+    return str(BgColors[name])
 
 
 ansi.fg_lookup = get_fg
@@ -84,8 +85,8 @@ class CmdLineApp(cmd2.Cmd):
     speak_parser.add_argument('-p', '--piglatin', action='store_true', help='atinLay')
     speak_parser.add_argument('-s', '--shout', action='store_true', help='N00B EMULATION MODE')
     speak_parser.add_argument('-r', '--repeat', type=int, help='output [n] times')
-    speak_parser.add_argument('-f', '--fg', choices=FG_COLORS, help='foreground color to apply to output')
-    speak_parser.add_argument('-b', '--bg', choices=BG_COLORS, help='background color to apply to output')
+    speak_parser.add_argument('-f', '--fg', choices=FgColors.colors(), help='foreground color to apply to output')
+    speak_parser.add_argument('-b', '--bg', choices=BgColors.colors(), help='background color to apply to output')
     speak_parser.add_argument('-l', '--bold', action='store_true', help='bold the output')
     speak_parser.add_argument('-u', '--underline', action='store_true', help='underline the output')
     speak_parser.add_argument('words', nargs='+', help='words to say')