33Support for ANSI escape sequences which are used for things like applying style to text,
44setting the window title, and asynchronous alerts.
55 """
6- from enum import Enum , unique
76import functools
87import re
8+ from enum import Enum
99from typing import Any , IO , List , Union
1010
1111import colorama
2727ANSI_STYLE_RE = re .compile (r'\x1b\[[^m]*m' )
2828
2929
30+ class ColorBase (Enum ):
31+ """
32+ Base class used for defining color enums. See fg and bg classes for examples.
33+
34+ Child classes should define enums in the follow structure:
35+ key: color name (e.g. black)
36+ value: anything that when cast to a string returns an ANSI sequence
37+ """
38+ def __str__ (self ) -> str :
39+ """
40+ Return ANSI color sequence instead of enum name
41+ This is helpful when using a ColorBase in an f-string or format() call
42+ e.g. my_str = "{}hello{}".format(fg.blue, fg.reset)
43+ """
44+ return str (self .value )
45+
46+ def __add__ (self , other : Any ) -> str :
47+ """
48+ Support building a color string when self is the left operand
49+ e.g. fg.blue + "hello"
50+ """
51+ return str (self ) + other
52+
53+ def __radd__ (self , other : Any ) -> str :
54+ """
55+ Support building a color string when self is the right operand
56+ e.g. "hello" + fg.reset
57+ """
58+ return other + str (self )
59+
60+ @classmethod
61+ def colors (cls ) -> List [str ]:
62+ """Return a list of color names."""
63+ # Use __members__ to ensure we get all key names, including those which are aliased
64+ return [color for color in cls .__members__ ]
65+
66+
3067# Foreground colors
31- # noinspection PyPep8Naming,DuplicatedCode
32- @unique
33- class fg (Enum ):
34- """Enum class for foreground colors (to support IDE autocompletion)."""
68+ # noinspection PyPep8Naming
69+ class fg (ColorBase ):
70+ """Enum class for foreground colors"""
3571 black = Fore .BLACK
3672 red = Fore .RED
3773 green = Fore .GREEN
@@ -50,26 +86,11 @@ class fg(Enum):
5086 bright_white = Fore .LIGHTWHITE_EX
5187 reset = Fore .RESET
5288
53- def __str__ (self ) -> str :
54- """Make the value the string representation instead of the enum name."""
55- return self .value
56-
57- @staticmethod
58- def colors () -> List [str ]:
59- """Return a list of color names."""
60- return [color .name for color in fg ]
61-
62- @staticmethod
63- def get_value (name : str ) -> str :
64- """Retrieve color code by name string."""
65- return fg .__members__ [name ].value
66-
6789
6890# Background colors
69- # noinspection PyPep8Naming,DuplicatedCode
70- @unique
71- class bg (Enum ):
72- """Enum class for background colors (to support IDE autocompletion)."""
91+ # noinspection PyPep8Naming
92+ class bg (ColorBase ):
93+ """Enum class for background colors"""
7394 black = Back .BLACK
7495 red = Back .RED
7596 green = Back .GREEN
@@ -88,20 +109,6 @@ class bg(Enum):
88109 bright_white = Back .LIGHTWHITE_EX
89110 reset = Back .RESET
90111
91- def __str__ (self ) -> str :
92- """Make the value the string representation instead of the enum name."""
93- return self .value
94-
95- @staticmethod
96- def colors () -> List [str ]:
97- """Return a list of color names."""
98- return [color .name for color in bg ]
99-
100- @staticmethod
101- def get_value (name : str ) -> str :
102- """Retrieve color code by name string."""
103- return bg .__members__ [name ].value
104-
105112
106113FG_RESET = fg .reset .value
107114BG_RESET = bg .reset .value
@@ -162,7 +169,7 @@ def fg_lookup(fg_name: Union[str, fg]) -> str:
162169 return fg_name .value
163170
164171 try :
165- ansi_escape = fg . get_value ( fg_name .lower ())
172+ ansi_escape = fg [ fg_name .lower ()]. value
166173 except KeyError :
167174 raise ValueError ('Foreground color {!r} does not exist; must be one of: {}' .format (fg_name , fg .colors ()))
168175 return ansi_escape
@@ -180,7 +187,7 @@ def bg_lookup(bg_name: Union[str, bg]) -> str:
180187 return bg_name .value
181188
182189 try :
183- ansi_escape = bg . get_value ( bg_name .lower ())
190+ ansi_escape = bg [ bg_name .lower ()]. value
184191 except KeyError :
185192 raise ValueError ('Background color {!r} does not exist; must be one of: {}' .format (bg_name , bg .colors ()))
186193 return ansi_escape
@@ -195,8 +202,10 @@ def style(text: Any, *, fg: Union[str, fg] = '', bg: Union[str, bg] = '', bold:
195202 to undo whatever styling was done at the beginning.
196203
197204 :param text: Any object compatible with str.format()
198- :param fg: foreground color. Relies on `fg_lookup()` to retrieve ANSI escape based on name or enum. Defaults to no color.
199- :param bg: background color. Relies on `bg_lookup()` to retrieve ANSI escape based on name or enum. Defaults to no color.
205+ :param fg: foreground color. Relies on `fg_lookup()` to retrieve ANSI escape based on name or enum.
206+ Defaults to no color.
207+ :param bg: background color. Relies on `bg_lookup()` to retrieve ANSI escape based on name or enum.
208+ Defaults to no color.
200209 :param bold: apply the bold style if True. Can be combined with dim. Defaults to False.
201210 :param dim: apply the dim style if True. Can be combined with bold. Defaults to False.
202211 :param underline: apply the underline style if True. Defaults to False.
0 commit comments