Merge pull request #876 from python-cmd2/color_autocomplete

Added convenience enums of fg and bg colors

9 files changed, 142 insertions, 79 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 77ba6b00..07225efc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,10 @@
 * Enhancements
     * Changed the default help text to make `help -v` more discoverable
     * Added `add_settable()` and `remove_settable()` convenience methods to update `self.settable` dictionary
+    * Added convenience `ansi.fg` and `ansi.bg` enums of foreground and background colors
+        * `ansi.style()` `fg` argument can now either be of type `str` or `ansi.fg`
+        * `ansi.style()` `bg` argument can now either be of type `str` or `ansi.bg`
+        * This supports IDE auto-completion of color names
 * Breaking changes
     * Renamed `locals_in_py` attribute of `cmd2.Cmd` to `self_in_py`
     * The following public attributes of `cmd2.Cmd` are no longer settable at runtime by default:
diff --git a/cmd2/__init__.py b/cmd2/__init__.py
index cc5a0963..43578e46 100644
--- a/cmd2/__init__.py
+++ b/cmd2/__init__.py
@@ -10,7 +10,7 @@ except DistributionNotFound:
     # package is not installed
     pass
 
-from .ansi import style
+from .ansi import style, fg, bg
 from .argparse_custom import Cmd2ArgumentParser, CompletionError, CompletionItem, set_default_argument_parser
 
 # Check if user has defined a module that sets a custom value for argparse_custom.DEFAULT_ARGUMENT_PARSER
diff --git a/cmd2/ansi.py b/cmd2/ansi.py
index d3af525a..abe7b3dc 100644
--- a/cmd2/ansi.py
+++ b/cmd2/ansi.py
@@ -3,9 +3,10 @@
 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 typing import Any, IO
+from typing import Any, IO, List, Union
 
 import colorama
 from colorama import Fore, Back, Style
@@ -25,57 +26,91 @@ allow_style = STYLE_TERMINAL
 # 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,
-}
-
-# 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,
-}
-
-FG_RESET = FG_COLORS['reset']
-BG_RESET = BG_COLORS['reset']
+
+# Foreground colors
+# noinspection PyPep8Naming,DuplicatedCode
+@unique
+class fg(Enum):
+    """Enum class for foreground colors (to support IDE autocompletion)."""
+    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
+
+    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]
+
+    @staticmethod
+    def get_value(name: str) -> str:
+        """Retrieve color code by name string."""
+        return fg.__members__[name].value
+
+
+# Background colors
+# noinspection PyPep8Naming,DuplicatedCode
+@unique
+class bg(Enum):
+    """Enum class for background colors (to support IDE autocompletion)."""
+    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
+
+    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(name: str) -> str:
+        """Retrieve color code by name string."""
+        return bg.__members__[name].value
+
+
+FG_RESET = fg.reset.value
+BG_RESET = bg.reset.value
 RESET_ALL = Style.RESET_ALL
 
 # Text intensities
 INTENSITY_BRIGHT = Style.BRIGHT
 INTENSITY_DIM = Style.DIM
 INTENSITY_NORMAL = Style.NORMAL
-
 # ANSI style sequences not provided by colorama
 UNDERLINE_ENABLE = colorama.ansi.code_to_chars(4)
 UNDERLINE_DISABLE = colorama.ansi.code_to_chars(24)
@@ -115,37 +150,44 @@ def style_aware_write(fileobj: IO, msg: str) -> None:
     fileobj.write(msg)
 
 
-def fg_lookup(fg_name: str) -> str:
+def fg_lookup(fg_name: Union[str, fg]) -> str:
     """
     Look up ANSI escape codes based on foreground color name.
 
-    :param fg_name: foreground color name to look up ANSI escape code(s) for
+    :param fg_name: foreground color name or enum to look up ANSI escape code(s) for
     :return: ANSI escape code(s) associated with this color
     :raises ValueError: if the color cannot be found
     """
+    if isinstance(fg_name, fg):
+        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))
+        raise ValueError('Foreground color {!r} does not exist; must be one of: {}'.format(fg_name, fg.colors()))
     return ansi_escape
 
 
-def bg_lookup(bg_name: str) -> str:
+def bg_lookup(bg_name: Union[str, bg]) -> str:
     """
     Look up ANSI escape codes based on background color name.
 
-    :param bg_name: background color name to look up ANSI escape code(s) for
+    :param bg_name: background color name or enum to look up ANSI escape code(s) for
     :return: ANSI escape code(s) associated with this color
     :raises ValueError: if the color cannot be found
     """
+    if isinstance(bg_name, bg):
+        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))
+        raise ValueError('Background color {!r} does not exist; must be one of: {}'.format(bg_name, bg.colors()))
     return ansi_escape
 
 
-def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False,
+# noinspection PyShadowingNames
+def style(text: Any, *, fg: Union[str, fg] = '', bg: Union[str, bg] = '', bold: bool = False,
           dim: bool = False, underline: bool = False) -> str:
     """
     Apply ANSI colors and/or styles to a string and return it.
@@ -153,8 +195,8 @@ def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False,
     to undo whatever styling was done at the beginning.
 
     :param text: Any object compatible with str.format()
-    :param fg: foreground color. Relies on `fg_lookup()` to retrieve ANSI escape based on name. Defaults to no color.
-    :param bg: background color. Relies on `bg_lookup()` to retrieve ANSI escape based on name. Defaults to no color.
+    :param fg: foreground color. Relies on `fg_lookup()` to retrieve ANSI escape based on name or enum. Defaults to no color.
+    :param bg: background color. Relies on `bg_lookup()` to retrieve ANSI escape based on name or enum. Defaults to no color.
     :param bold: apply the bold style if True. Can be combined with dim. Defaults to False.
     :param dim: apply the dim style if True. Can be combined with bold. Defaults to False.
     :param underline: apply the underline style if True. Defaults to False.
@@ -197,13 +239,13 @@ def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False,
 # Default styles for printing strings of various types.
 # These can be altered to suit an application's needs and only need to be a
 # function with the following structure: func(str) -> str
-style_success = functools.partial(style, fg='green')
+style_success = functools.partial(style, fg=fg.green)
 """Partial function supplying arguments to :meth:`cmd2.ansi.style()` which colors text to signify success"""
 
-style_warning = functools.partial(style, fg='bright_yellow')
+style_warning = functools.partial(style, fg=fg.bright_yellow)
 """Partial function supplying arguments to :meth:`cmd2.ansi.style()` which colors text to signify a warning"""
 
-style_error = functools.partial(style, fg='bright_red')
+style_error = functools.partial(style, fg=fg.bright_red)
 """Partial function supplying arguments to :meth:`cmd2.ansi.style()` which colors text to signify an error"""
 
 
diff --git a/docs/features/initialization.rst b/docs/features/initialization.rst
index 5721caa3..6824c7bf 100644
--- a/docs/features/initialization.rst
+++ b/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, _):
diff --git a/examples/basic.py b/examples/basic.py
index 37db6501..5df3d1b5 100755
--- a/examples/basic.py
+++ b/examples/basic.py
@@ -9,7 +9,7 @@
     6) Shell-like capabilities
 """
 import cmd2
-from cmd2 import style
+from cmd2 import style, fg, bg
 
 
 class BasicApp(cmd2.Cmd):
@@ -19,7 +19,7 @@ class BasicApp(cmd2.Cmd):
         super().__init__(multiline_commands=['echo'], persistent_history_file='cmd2_history.dat',
                          startup_script='scripts/startup.txt', use_ipython=True)
 
-        self.intro = style('Welcome to PyOhio 2019 and cmd2!', fg='red', bg='white', bold=True) + ' 😀'
+        self.intro = style('Welcome to PyOhio 2019 and cmd2!', fg=fg.red, bg=bg.white, bold=True) + ' 😀'
 
         # Allow access to your application in py and ipy via self
         self.self_in_py = True
diff --git a/examples/colors.py b/examples/colors.py
index 33b17e53..1466471b 100755
--- a/examples/colors.py
+++ b/examples/colors.py
@@ -48,8 +48,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=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')
diff --git a/examples/initialization.py b/examples/initialization.py
index c13ed137..609255f1 100755
--- a/examples/initialization.py
+++ b/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 @@ class BasicApp(cmd2.Cmd):
                          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 @@ class BasicApp(cmd2.Cmd):
 
         # 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, _):
diff --git a/examples/pirate.py b/examples/pirate.py
index acbab17c..a50f9a51 100755
--- a/examples/pirate.py
+++ b/examples/pirate.py
@@ -25,7 +25,7 @@ class Pirate(cmd2.Cmd):
         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
diff --git a/tests/test_ansi.py b/tests/test_ansi.py
index cb68cb28..51c3b50f 100644
--- a/tests/test_ansi.py
+++ b/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():
@@ -119,3 +119,22 @@ def test_async_alert_str(cols, prompt, line, cursor, msg, expected):
     alert_str = ansi.async_alert_str(terminal_columns=cols, prompt=prompt, line=line, cursor_offset=cursor,
                                      alert_msg=msg)
     assert alert_str == expected
+
+
+def test_fg_enum():
+    assert ansi.fg_lookup('bright_red') == ansi.fg_lookup(ansi.fg.bright_red)
+
+def test_fg_enum_to_str():
+    assert str(ansi.fg.black) == ansi.fg_lookup('black')
+
+def test_bg_enum():
+    assert ansi.bg_lookup('green') == ansi.bg_lookup(ansi.bg.green)
+
+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()