Having two parallel datastructures each for foreground and background colors felt really clunky - now we just have one

The Enum classes are now smart and deal with it all.

Author: tleonhardt

cmd2/ansi.py

@@ -6,7 +6,7 @@
 from enum import Enum, unique
 import functools
 import re
-from typing import Any, IO, Union
+from typing import Any, IO, List, Union
 
 import colorama
 from colorama import Fore, Back, Style
@@ -26,28 +26,7 @@
 # Regular expression to match ANSI style sequences (including 8-bit and 24-bit colors)
 ANSI_STYLE_RE = re.compile(r'\x1b\[[^m]*m')
 
-# Foreground color presets
-FG_COLORS = {
-    'black': Fore.BLACK,
-    'red': Fore.RED,
-    'green': Fore.GREEN,
-    'yellow': Fore.YELLOW,
-    'blue': Fore.BLUE,
-    'magenta': Fore.MAGENTA,
-    'cyan': Fore.CYAN,
-    'white': Fore.WHITE,
-    'bright_black': Fore.LIGHTBLACK_EX,
-    'bright_red': Fore.LIGHTRED_EX,
-    'bright_green': Fore.LIGHTGREEN_EX,
-    'bright_yellow': Fore.LIGHTYELLOW_EX,
-    'bright_blue': Fore.LIGHTBLUE_EX,
-    'bright_magenta': Fore.LIGHTMAGENTA_EX,
-    'bright_cyan': Fore.LIGHTCYAN_EX,
-    'bright_white': Fore.LIGHTWHITE_EX,
-    'reset': Fore.RESET,
-}
-
-
+# Foreground colors
 @unique
 class fg(Enum):
     """Enum class for foreground colors (to support IDE autocompletion)."""
@@ -69,33 +48,22 @@ class fg(Enum):
     bright_white = Fore.LIGHTWHITE_EX
     reset = Fore.RESET
 
-    def __str__(self):
+    def __str__(self) -> str:
         """Make the value the string representation instead of the enum name."""
         return self.value
 
+    @staticmethod
+    def colors() -> List[str]:
+        """Return a list of color names."""
+        return [color.name for color in fg]
 
-# Background color presets
-BG_COLORS = {
-    'black': Back.BLACK,
-    'red': Back.RED,
-    'green': Back.GREEN,
-    'yellow': Back.YELLOW,
-    'blue': Back.BLUE,
-    'magenta': Back.MAGENTA,
-    'cyan': Back.CYAN,
-    'white': Back.WHITE,
-    'bright_black': Back.LIGHTBLACK_EX,
-    'bright_red': Back.LIGHTRED_EX,
-    'bright_green': Back.LIGHTGREEN_EX,
-    'bright_yellow': Back.LIGHTYELLOW_EX,
-    'bright_blue': Back.LIGHTBLUE_EX,
-    'bright_magenta': Back.LIGHTMAGENTA_EX,
-    'bright_cyan': Back.LIGHTCYAN_EX,
-    'bright_white': Back.LIGHTWHITE_EX,
-    'reset': Back.RESET,
-}
+    @staticmethod
+    def get_value(color: str) -> str:
+        """Retrieve color code by name string."""
+        return fg.__members__[color].value
 
 
+# Background colors
 @unique
 class bg(Enum):
     """Enum class for background colors (to support IDE autocompletion)."""
@@ -117,13 +85,23 @@ class bg(Enum):
     bright_white = Back.LIGHTWHITE_EX
     reset = Back.RESET
 
-    def __str__(self):
+    def __str__(self) -> str:
         """Make the value the string representation instead of the enum name."""
         return self.value
 
+    @staticmethod
+    def colors() -> List[str]:
+        """Return a list of color names."""
+        return [color.name for color in bg]
+
+    @staticmethod
+    def get_value(color: str) -> str:
+        """Retrieve color code by name string."""
+        return bg.__members__[color].value
+
 
-FG_RESET = FG_COLORS['reset']
-BG_RESET = BG_COLORS['reset']
+FG_RESET = fg.reset.value
+BG_RESET = bg.reset.value
 RESET_ALL = Style.RESET_ALL
 
 # Text intensities
@@ -181,7 +159,7 @@ def fg_lookup(fg_name: Union[str, fg]) -> str:
         return fg_name.value
 
     try:
-        ansi_escape = FG_COLORS[fg_name.lower()]
+        ansi_escape = fg.get_value(fg_name.lower())
     except KeyError:
         raise ValueError('Foreground color {!r} does not exist.'.format(fg_name))
     return ansi_escape
@@ -199,7 +177,7 @@ def bg_lookup(bg_name: Union[str, bg]) -> str:
         return bg_name.value
 
     try:
-        ansi_escape = BG_COLORS[bg_name.lower()]
+        ansi_escape = bg.get_value(bg_name.lower())
     except KeyError:
         raise ValueError('Background color {!r} does not exist.'.format(bg_name))
     return ansi_escape

docs/features/initialization.rst

@@ -19,8 +19,7 @@ capabilities which you may wish to utilize while initializing the app::
         10) How to make custom attributes settable at runtime
     """
     import cmd2
-    from cmd2 import style
-    from cmd2.ansi import FG_COLORS
+    from cmd2 import style, fg, bg
 
     class BasicApp(cmd2.Cmd):
         CUSTOM_CATEGORY = 'My Custom Commands'
@@ -30,7 +29,7 @@ capabilities which you may wish to utilize while initializing the app::
                              startup_script='scripts/startup.txt', use_ipython=True)
 
             # Prints an intro banner once upon application startup
-            self.intro = style('Welcome to cmd2!', fg='red', bg='white', bold=True)
+            self.intro = style('Welcome to cmd2!', fg=fg.red, bg=bg.white, bold=True)
 
             # Show this as the prompt when asking for input
             self.prompt = 'myapp> '
@@ -51,7 +50,7 @@ capabilities which you may wish to utilize while initializing the app::
             self.add_settable(cmd2.Settable('foreground_color',
                                             str,
                                             'Foreground color to use with echo command',
-                                            choices=FG_COLORS))
+                                            choices=fg.colors()))
 
         @cmd2.with_category(CUSTOM_CATEGORY)
         def do_intro(self, _):

examples/colors.py

@@ -48,8 +48,8 @@ def __init__(self):
     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=ansi.FG_COLORS, help='foreground color to apply to output')
-    speak_parser.add_argument('-b', '--bg', choices=ansi.BG_COLORS, help='background color to apply to output')
+    speak_parser.add_argument('-f', '--fg', choices=ansi.fg.colors(), help='foreground color to apply to output')
+    speak_parser.add_argument('-b', '--bg', choices=ansi.bg.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')

examples/initialization.py

@@ -13,8 +13,7 @@
     10) How to make custom attributes settable at runtime
 """
 import cmd2
-from cmd2 import style
-from cmd2.ansi import FG_COLORS
+from cmd2 import style, fg, bg
 
 
 class BasicApp(cmd2.Cmd):
@@ -25,7 +24,7 @@ def __init__(self):
                          startup_script='scripts/startup.txt', use_ipython=True)
 
         # Prints an intro banner once upon application startup
-        self.intro = style('Welcome to cmd2!', fg='red', bg='white', bold=True)
+        self.intro = style('Welcome to cmd2!', fg=fg.red, bg=bg.white, bold=True)
 
         # Show this as the prompt when asking for input
         self.prompt = 'myapp> '
@@ -44,7 +43,7 @@ def __init__(self):
 
         # Make echo_fg settable at runtime
         self.add_settable(cmd2.Settable('foreground_color', str, 'Foreground color to use with echo command',
-                                        choices=FG_COLORS))
+                                        choices=fg.colors()))
 
     @cmd2.with_category(CUSTOM_CATEGORY)
     def do_intro(self, _):

examples/pirate.py

@@ -25,7 +25,7 @@ def __init__(self):
         self.songcolor = 'blue'
 
         # Make songcolor settable at runtime
-        self.add_settable(cmd2.Settable('songcolor', str, 'Color to ``sing``', choices=cmd2.ansi.FG_COLORS))
+        self.add_settable(cmd2.Settable('songcolor', str, 'Color to ``sing``', choices=cmd2.ansi.fg.colors()))
 
         # prompts and defaults
         self.gold = 0

tests/test_ansi.py

@@ -32,14 +32,14 @@ def test_style_none():
 def test_style_fg():
     base_str = HELLO_WORLD
     fg_color = 'blue'
-    ansi_str = ansi.FG_COLORS[fg_color] + base_str + ansi.FG_RESET
+    ansi_str = ansi.fg.get_value(fg_color) + base_str + ansi.FG_RESET
     assert ansi.style(base_str, fg=fg_color) == ansi_str
 
 
 def test_style_bg():
     base_str = HELLO_WORLD
     bg_color = 'green'
-    ansi_str = ansi.BG_COLORS[bg_color] + base_str + ansi.BG_RESET
+    ansi_str = ansi.bg.get_value(bg_color) + base_str + ansi.BG_RESET
     assert ansi.style(base_str, bg=bg_color) == ansi_str
 
 
@@ -65,7 +65,7 @@ def test_style_multi():
     base_str = HELLO_WORLD
     fg_color = 'blue'
     bg_color = 'green'
-    ansi_str = (ansi.FG_COLORS[fg_color] + ansi.BG_COLORS[bg_color] +
+    ansi_str = (ansi.fg.get_value(fg_color) + ansi.bg.get_value(bg_color) +
                 ansi.INTENSITY_BRIGHT + ansi.INTENSITY_DIM + ansi.UNDERLINE_ENABLE +
                 base_str +
                 ansi.FG_RESET + ansi.BG_RESET +
@@ -85,7 +85,7 @@ def test_style_color_not_exist():
 
 def test_fg_lookup_exist():
     fg_color = 'green'
-    assert ansi.fg_lookup(fg_color) == ansi.FG_COLORS[fg_color]
+    assert ansi.fg_lookup(fg_color) == ansi.fg.get_value(fg_color)
 
 
 def test_fg_lookup_nonexist():
@@ -95,7 +95,7 @@ def test_fg_lookup_nonexist():
 
 def test_bg_lookup_exist():
     bg_color = 'green'
-    assert ansi.bg_lookup(bg_color) == ansi.BG_COLORS[bg_color]
+    assert ansi.bg_lookup(bg_color) == ansi.bg.get_value(bg_color)
 
 
 def test_bg_lookup_nonexist():
@@ -133,3 +133,8 @@ def test_bg_enum():
 def test_bg_enum_to_str():
     assert str(ansi.bg.blue) == ansi.bg_lookup('blue')
 
+def test_fg_colors():
+    assert list(ansi.fg.__members__.keys()) == ansi.fg.colors()
+
+def test_bg_colors():
+    assert list(ansi.bg.__members__.keys()) == ansi.bg.colors()