Merge pull request #876 from python-cmd2/color_autocomplete

Added convenience enums of fg and bg colors

Author: tleonhardt

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:

cmd2/__init__.py

@@ -10,7 +10,7 @@
     # 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

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 @@
 # 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,46 +150,53 @@ 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.
     The styling is self contained which means that at the end of the string reset code(s) are issued
     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"""
 
 

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/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 @@ def __init__(self):
         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

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