Added convenience ansi.fg and ansi.bg enums of foreground and background colors which style() can now optionally use

This is to make it easier to autocomplete color names in an IDE

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, Union
 
 import colorama
 from colorama import Fore, Back, Style
@@ -46,6 +47,33 @@
     'reset': Fore.RESET,
 }
 
+
+@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):
+        """Make the value the string representation instead of the enum name."""
+        return self.value
+
+
 # Background color presets
 BG_COLORS = {
     'black': Back.BLACK,
@@ -67,6 +95,33 @@
     'reset': Back.RESET,
 }
 
+
+@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):
+        """Make the value the string representation instead of the enum name."""
+        return self.value
+
+
 FG_RESET = FG_COLORS['reset']
 BG_RESET = BG_COLORS['reset']
 RESET_ALL = Style.RESET_ALL
@@ -75,7 +130,6 @@
 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 +169,52 @@ 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()]
     except KeyError:
         raise ValueError('Foreground color {!r} does not exist.'.format(fg_name))
     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()]
     except KeyError:
         raise ValueError('Background color {!r} does not exist.'.format(bg_name))
     return ansi_escape
 
 
-def style(text: Any, *, fg: str = '', bg: str = '', bold: bool = False,
+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 +257,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"""
 
 

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

tests/test_ansi.py

@@ -119,3 +119,17 @@ 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')
+