jarvis-core/jarvis_core/util/ansi.py

75 lines
1.4 KiB
Python

"""ANSI color helper."""
from enum import Enum, EnumMeta
from typing import Any, List
ESCAPE = "\u001b"
RESET = ESCAPE + "[0m"
class SearchableEnum(EnumMeta):
def __contains__(cls, item: Any):
return item in [v.value for v in cls.__members__.values()]
class Format(Enum, metaclass=SearchableEnum):
NORMAL = 0
BOLD = 1
UNDERLINE = 4
class Fore(Enum, metaclass=SearchableEnum):
GRAY = 30
RED = 31
GREEN = 32
YELLOW = 33
BLUE = 34
PINK = 35
CYAN = 36
WHITE = 37
class Back(Enum, metaclass=SearchableEnum):
DARK_BLUE = 40
ORANGE = 41
DARK_GRAY = 42
GRAY = 43
LIGHT_GRAY = 44
INDIGO = 45
LIGHTER_GRAY = 46
WHITE = 47
def fmt(*formats: List[Format | Fore | Back] | int) -> str:
"""
Get format string.
Args:
formats: List of format modifiers
Return:
Format string
"""
fore = ""
back = ""
fmt = ""
block_fmt = False
for f in formats:
if isinstance(f, Enum):
f = f.value
if f == 0:
fmt = "0;"
elif f in [1, 4]:
if str(f) not in fmt and not block_fmt:
fmt += f"{f};"
elif 30 <= f <= 37:
fore = f"{f};"
elif 40 <= f <= 47:
back = f"{f};"
ret = fmt + fore + back
if not any(ret, fore, back):
ret = RESET
if ret[-1] == ";":
ret = ret[:-1]
return ESCAPE + "[" + ret + "m"