imp/sandbox/cli/style.py

'''
=== ANSI Escape Sequences ===
This file exists to organise and name
all important ANSI escape sequences
for management of the CLI.
'''

from sys import stdout
from enum import StrEnum

# Removes ALL set colors/styling
RESET = '\x1b[0m'

'''
Implements text foreground coloring. 
'''
class Color(StrEnum):
    BLACK   = '\x1b[30m'
    RED     = '\x1b[31m'
    GREEN   = '\x1b[32m'
    YELLOW  = '\x1b[33m'
    BLUE    = '\x1b[34m'
    MAGENTA = '\x1b[35m'
    CYAN    = '\x1b[36m'
    WHITE   = '\x1b[37m'

    '''
    Handles the application of a Color to a given text.
    Unset the `temp` flag to leave this style on (after the given text) 
    '''
    @staticmethod
    def apply(color: 'Color', text: str, temp: bool = True) -> str:
        return color + text + Color.WHITE if temp else ''

'''
Implements text effects.  
'''
class Effect(StrEnum):
    BOLD      = '\x1b[1m'
    DIM       = '\x1b[2m'
    ITALICS   = '\x1b[3m'
    UNDERLINE = '\x1b[4m'
    BLINK     = '\x1b[5m'
    REVERSE   = '\x1b[7m'
    HIDE      = '\x1b[8m'

    '''
    Returns the opposite ESC code to a given Effect.
    ie the opposite of BOLD '\x1b[1m]' is '\x1b[21m'
    '''
    @staticmethod
    def _inverse(effect: 'Effect') -> str:
        return f'{effect[:2]}2{effect[2:]}'

    '''
    Handles the application of a Effect to a given text.
    Unset the `temp` flag to leave this style on (after the given text) 
    '''
    @staticmethod
    def apply(effect: 'Effect', text: str, temp: bool = True) -> str:
        return effect + text + Effect._inverse(effect) if temp else ''

'''
Applies styling (color/effects/etc) to a given string
Unset the `temp` flag to leave this style on (after the given text) 
'''
def style(text: str, *args: Color | Effect, temp: bool = True) -> str:
    # unsures we don't add redundant "set white" commands
    unset_color = False
    for arg in args:
        unset = temp
        if isinstance(arg, Color):
            unset_color |= temp
            unset = False
        text = type(arg).apply(arg, text, temp=True)
    if unset_color:
        text += Color.WHITE
    return text


'''
Implements cursor movement functionality.
NOTE: 
    The Cursor class currently has no ability
    to handle EXACT line (row) numbers. Currently
    I'm assuming all functionality of noether can
    be implemented solely via relative movements.
'''
class Cursor(StrEnum):
    # SAVE current / RESTORE last saved cursor position
    _SAVE = f'\x1b[7'
    _RESTORE = f'\x1b[8'
    
    _MV_UP     = 'A'
    _MV_DOWN   = 'B'
    _MV_LEFT   = 'C'
    _MV_RIGHT  = 'D'
    # NEXT/PREV are the same as DOWN/UP (respectively)
    # except that the cursor will reset to column 0
    _MV_NEXT   = 'E' 
    _MV_PREV   = 'F'
    # move cursor to the start of the current line
    _MV_START  = '\r' # there is no ESC CSI for carriage return
    
    _MV_COLUMN = 'G'
    
    '''
    Generates an ESC code sequence corresponding
    to a horizontal movement relative to the 
    cursor's current vertical position.
        +ive = right
        -ive = left   
    '''
    @staticmethod
    def _XMOVE_REL(n: int) -> str:
        if n > 0:
            return f'\x1b[{n}{Cursor._MV_RIGHT}'
        if n < 0:
            return f'\x1b[{n}{Cursor._MV_LEFT}'
        return ''

    '''
    Generates an ESC code sequence corresponding
    to a horizontal movement to an EXACT column.  
    '''
    @staticmethod
    def _XMOVE(n: int) -> str:
        return f'\x1b[{n}{Cursor._MV_COLUMN}'

    '''
    Generates an ESC code sequence corresponding
    to a vertical movement relative to the 
    cursor's current vertical position.
        +ive = down
        -ive = up      
    '''
    @staticmethod
    def _YMOVE_REL(n: int, reset: bool = False) -> str:
        if n > 0:
            if reset:
                return f'\x1b[{n}{Cursor._MV_NEXT}'
            return f'\x1b[{n}{Cursor._MV_DOWN}'
        if n < 0:
            if reset:
                return f'\x1b[{n}{Cursor._MV_PREV}'
            return f'\x1b[{n}{Cursor._MV_UP}'
        return ''

    '''
    Sets the cursor column (horizontal) position to an exact value.
    NOTE: does NOT flush stdout buffer
    '''
    @staticmethod
    def set_x(n: int) -> None:
        stdout.write(Cursor.XMOVE(n))

    '''
    Moves the cursor left/right n columns (relative).
    NOTE: does NOT flush stdout buffer
    '''
    @staticmethod
    def move_x(n: int) -> None:
        stdout.write(Cursor._XMOVE_REL(n))
    
    '''
    Moves the cursor up/down n rows and resets the
    cursor to be at the start of the line
    NOTE: does NOT flush stdout buffer
    '''
    @staticmethod
    def move_y(n: int, reset: bool = True) -> None:
        stdout.write(Cursor._YMOVE_REL(n, reset=reset))

    '''
    Saves the current cursor position.
    NOTE: does NOT flush stdout buffer
    '''
    @staticmethod
    def save() -> None:
        stdout.write(Cursor._SAVE)

    '''
    Restores the cursor position to a saved position.
    NOTE: does NOT flush stdout buffer
    '''
    @staticmethod
    def restore() -> None:
        stdout.write(Cursor._RESTORE)

'''
Handles erasing content displayed on the screen.
NOTE that the cursor position is NOT updated
via these sequences. The \r code should be given after.
'''
class Erase(StrEnum):
    # erase everything from the current cursor
    # position to the START/END of the screen
    _SCR_AFTER  = '\x1b[0J'
    _SCR_BEFORE = '\x1b[1J'
    _SCR_ALL    = '\x1b[3J' # erase screen and delete all saved cursors
    
    # erase everything from the current cursor
    # position to the START/END of the current line
    _LINE_AFTER  = '\x1b[0K'
    _LINE_BEFORE = '\x1b[1K'
    _LINE_ALL    = '\x1b[2K'

    '''
    Erase characters on the entire screen.
    Set `before` flag to only erase before the cursor,
    set `after` flag to only erase after the cursor.
    '''
    @staticmethod
    def screen(before: bool = False, after: bool = False) -> None:
        if before:
            stdout.write(Erase._SCR_BEFORE)
        elif after:
            stdout.write(Erase._SCR_AFTER)
        else:
            stdout.write(Erase._SCR_ALL)

    '''
    Erase characters on the current line.
    Set `before` flag to only erase before the cursor,
    set `after` flag to only erase after the cursor.
    '''
    @staticmethod
    def line(before: bool = False, after: bool = False) -> None:
        if before:
            stdout.write(Erase._LINE_BEFORE)
        elif after:
            stdout.write(Erase._LINE_AFTER)
        else:
            stdout.write(Erase._LINE_ALL)