diff options
Diffstat (limited to 'cmd2/ansi.py')
| -rw-r--r-- | cmd2/ansi.py | 84 |
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""" |