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

1 files changed, 72 insertions, 12 deletions

diff --git a/cmd2/ansi.py b/cmd2/ansi.py
index d3af525a..39175d3f 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, Union
 
 import colorama
 from colorama import Fore, Back, Style
@@ -46,6 +47,33 @@ FG_COLORS = {
     '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 @@ BG_COLORS = {
     '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 @@ RESET_ALL = Style.RESET_ALL
 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,14 +169,17 @@ 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:
@@ -130,14 +187,17 @@ def fg_lookup(fg_name: str) -> str:
     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:
@@ -145,7 +205,7 @@ def bg_lookup(bg_name: str) -> str:
     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.
@@ -153,8 +213,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 +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"""