noether/prbraid/color.py

from enum import Enum
from typing import Optional

class Color(Enum):
    Reset = 0
    Bold = 1
    Italics = 2  
    Underline = 4 
    
    Black = 30
    Red = 31
    Green = 32
    Yellow = 33
    Blue = 34
    Magenta = 35
    Cyan = 36
    White = 37

    # escape sequence format string
    __fescseq = '\033[{0}'
    # ansi reset code
    __ansi_rst = '\033[0m'
    
    @staticmethod
    def _ansi_ret(ansi):
        return ansi, Color.__ansi_white

    @staticmethod
    def code(color: 'Color'):
        return Color.__fescseq.format(f'{color.value}m')

    # move the cursor up n lines
    @staticmethod
    def mvup_code(n: int):
        return Color.__fescseq.format(f'{n}A')
    
    def ansi(self):
        code = Color.code(self)
        if self == Color.White:
            return code, ''
        return code, Color.__ansi_rst

# !! ULTRA PRINT !!
def uprint(text: str, 
           color: Optional[Color] = None,
           style: Optional[Color] = None,
           moveup: int = 0, 
           end: str = '\n', 
           flush: bool = True):
    if color is not None:
        c = color.ansi()
        text = f'{c[0]}{text}{c[1]}'
    if style is not None:
        s = style.ansi()
        tail = '' if (c is not None and c[1] == '') else s[1]
        text = f'{s[0]}{text}{tail}'
    if moveup > 0:
        text = Color.mvup_code(moveup) + text
    print(text, end=end, flush=flush)