diff options
Diffstat (limited to 'cmd2/ansi.py')
| -rw-r--r-- | cmd2/ansi.py | 91 |
1 files changed, 50 insertions, 41 deletions
diff --git a/cmd2/ansi.py b/cmd2/ansi.py index abe7b3dc..2813d62f 100644 --- a/cmd2/ansi.py +++ b/cmd2/ansi.py @@ -3,9 +3,9 @@ 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 enum import Enum from typing import Any, IO, List, Union import colorama @@ -27,11 +27,47 @@ allow_style = STYLE_TERMINAL ANSI_STYLE_RE = re.compile(r'\x1b\[[^m]*m') +class ColorBase(Enum): + """ + Base class used for defining color enums. See fg and bg classes for examples. + + Child classes should define enums in the follow structure: + key: color name (e.g. black) + value: anything that when cast to a string returns an ANSI sequence + """ + def __str__(self) -> str: + """ + Return ANSI color sequence instead of enum name + This is helpful when using a ColorBase in an f-string or format() call + e.g. my_str = "{}hello{}".format(fg.blue, fg.reset) + """ + return str(self.value) + + def __add__(self, other: Any) -> str: + """ + Support building a color string when self is the left operand + e.g. fg.blue + "hello" + """ + return str(self) + other + + def __radd__(self, other: Any) -> str: + """ + Support building a color string when self is the right operand + e.g. "hello" + fg.reset + """ + return other + str(self) + + @classmethod + def colors(cls) -> List[str]: + """Return a list of color names.""" + # Use __members__ to ensure we get all key names, including those which are aliased + return [color for color in cls.__members__] + + # Foreground colors -# noinspection PyPep8Naming,DuplicatedCode -@unique -class fg(Enum): - """Enum class for foreground colors (to support IDE autocompletion).""" +# noinspection PyPep8Naming +class fg(ColorBase): + """Enum class for foreground colors""" black = Fore.BLACK red = Fore.RED green = Fore.GREEN @@ -50,26 +86,11 @@ class fg(Enum): 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).""" +# noinspection PyPep8Naming +class bg(ColorBase): + """Enum class for background colors""" black = Back.BLACK red = Back.RED green = Back.GREEN @@ -88,20 +109,6 @@ class bg(Enum): 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 @@ -162,7 +169,7 @@ def fg_lookup(fg_name: Union[str, fg]) -> str: return fg_name.value try: - ansi_escape = fg.get_value(fg_name.lower()) + ansi_escape = fg[fg_name.lower()].value except KeyError: raise ValueError('Foreground color {!r} does not exist; must be one of: {}'.format(fg_name, fg.colors())) return ansi_escape @@ -180,7 +187,7 @@ def bg_lookup(bg_name: Union[str, bg]) -> str: return bg_name.value try: - ansi_escape = bg.get_value(bg_name.lower()) + ansi_escape = bg[bg_name.lower()].value except KeyError: raise ValueError('Background color {!r} does not exist; must be one of: {}'.format(bg_name, bg.colors())) return ansi_escape @@ -195,8 +202,10 @@ def style(text: Any, *, fg: Union[str, fg] = '', bg: Union[str, bg] = '', bold: 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 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 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. |