pyansi/pyansi/__init__.py

from base64 import b64encode
from dataclasses import dataclass
from enum import Enum
from enum import IntEnum
from enum import StrEnum
from enum import auto
from functools import wraps
import re
from typing import Annotated
from typing import Callable
from typing import Concatenate
from typing import ParamSpec
from typing import Self
from typing import get_args
from typing import get_origin
from typing import get_type_hints


@dataclass
class ValueRange:
    low: int
    high: int


@dataclass
class RegexMatch:
    pattern: re.Pattern


def check_annotations(func: Callable) -> Callable:
    @wraps(func)
    def wrapped(*args, **kwargs):
        type_hints = get_type_hints(func, include_extras=True)
        for param, hint in type_hints.items():
            if get_origin(hint) is not Annotated:
                continue
            hint_type, *hint_args = get_args(hint)
            if hint_type is list or get_origin(hint_type) is list:
                for arg in hint_args:
                    match arg:
                        case ValueRange():
                            value = kwargs[param]
                            if value < arg.low or value > arg.high:
                                raise ValueError(
                                    f"Parameter '{param}' must be >= {arg.low} and <= {arg.high}"
                                )
                        case RegexMatch():
                            value = kwargs[param]
                            if not arg.pattern.match(value):
                                raise ValueError(
                                    f"Parameter '{param}' must match pattern '{arg.pattern.pattern}'"
                                )

        return func(*args, **kwargs)

    return wrapped


class Mode(Enum):
    C0 = auto()
    C1 = auto()
    Bash = auto()


class Reset:
    """reset all graphics renditions to the default"""

    def __new__(cls: type[Self]) -> Self:
        if not hasattr(cls, "instance"):
            cls.instance = super(Reset, cls).__new__(cls)
        return cls.instance


class Color8(IntEnum):
    """8 color palette"""

    BLACK = 0
    RED = 1
    GREEN = 2
    YELLOW = 3
    BLUE = 4
    MAGENTA = 5
    CYAN = 6
    WHITE = 7


class Style(IntEnum):
    """enable graphics rendition style"""

    BOLD = 1
    DIM = 2
    FAINT = 2
    ITALIC = 3
    UNDERLINE = 4
    BLINKING = 5
    INVERSE = 7
    REVERSE = 7
    HIDDEN = 8
    INVISIBLE = 8
    STRIKETHROUGH = 9
    DOUBLE_UNDERLINE = 21


class ResetStyle(IntEnum):
    """reset graphics rendition style"""

    BOLD = 22
    DIM = 22
    FAINT = 22
    ITALIC = 23
    UNDERLINE = 24
    BLINKING = 25
    INVERSE = 27
    REVERSE = 27
    HIDDEN = 28
    INVISIBLE = 28
    STRIKETHROUGH = 29
    DOUBLE_UNDERLINE = 24


class UnderlineStyle(IntEnum):
    NO_UNDERLINE = 0
    STRAIGHT_UNDERLINE = 1
    DOUBLE_UNDERLINE = 2
    CURLY_UNDERLINE = 3
    DOTTED_UNDERLINE = 4
    DASHED_UNDERLINE = 5


@dataclass
class Color256:
    color: int


@dataclass
class ColorRGB:
    red: int
    green: int
    blue: int


class Layer(Enum):
    FOREGROUND = auto()
    BACKGROUND = auto()
    UNDERLINE = auto()


@dataclass
class Color:
    color: Color8 | Color256 | ColorRGB
    layer: Layer = Layer.FOREGROUND
    bright: bool = False


GraphicsRendtion = Reset | Style | ResetStyle | UnderlineStyle | Color


class EraseDisplayMode(IntEnum):
    TO_END_OF_SCREEN = 0
    TO_CURSOR = 1
    ENTIRE_SCREEN = 2
    SAVED_LINES = 3


class EraseLineMode(IntEnum):
    TO_END_OF_LINE = 0
    TO_CURSOR = 1
    ENTIRE_LINE = 2


class ActiveStatusDisplayMode(IntEnum):
    MAIN_DISPLAY = 0
    STATUS_LINE = 1


class StatusLineType(IntEnum):
    NONE = 0
    INIDICATOR = 1
    HOST_WRITABLE = 2


class DesktopNotificationPart(StrEnum):
    TITLE = "title"
    BODY = "body"


class DesktopNotificationDone(IntEnum):
    NOT_DONE = 0
    DONE = 1


class ClipboardType(StrEnum):
    CLIPBOARD = "c"
    PRIMARY = "p"
    SELECT = "s"
    CUT_BUFFER_0 = "0"
    CUT_BUFFER_1 = "1"
    CUT_BUFFER_2 = "2"
    CUT_BUFFER_3 = "3"
    CUT_BUFFER_4 = "4"
    CUT_BUFFER_5 = "5"
    CUT_BUFFER_6 = "6"
    CUT_BUFFER_7 = "7"
    CUT_BUFFER_8 = "8"


class CursorShape(IntEnum):
    BLOCK = 0
    VERTICAL_BAR = 1
    UNDERLINE = 2


P = ParamSpec("P")


class ANSI:
    mode: Mode
    wrap: bool

    @staticmethod
    def readline_wrap(
        func: Callable[Concatenate["ANSI", P], str]
    ) -> Callable[Concatenate["ANSI", P], str]:
        """wrap output in SOH/STX so that readline handles prompts properly"""

        @wraps(func)
        def wrapped(ansi: ANSI, *args: P.args, **kwargs: P.kwargs) -> str:
            result = ""
            if ansi.wrap:
                result += ansi.SOH
            result += func(ansi, *args, **kwargs)
            if ansi.wrap:
                result += ansi.STX
            return result

        return wrapped

    def __init__(self: Self, mode: Mode = Mode.C0, wrap: bool = False) -> None:
        self.mode = mode
        self.wrap = wrap

    def _wrap(self: Self, data: str) -> str:
        if self.wrap:
            return f"{self.SOH}{data}{self.STX}"
        return data

    NUL: Annotated[str, "null"] = "\x00"

    @property
    def SOH(self: Self) -> str:
        """start of heading"""
        match self.mode:
            case Mode.C0 | Mode.C1:
                return "\x01"
            case Mode.Bash:
                return "\\["

    @property
    def STX(self: Self) -> str:
        """start of text"""
        match self.mode:
            case Mode.C0 | Mode.C1:
                return "\x02"
            case Mode.Bash:
                return "\\]"

    @property
    def RL_PROMPT_START_IGNORE(self: Self) -> str:
        if self.wrap:
            return self.SOH
        return ""

    @property
    def RL_PROMPT_END_IGNORE(self: Self) -> str:
        if self.wrap:
            return self.STX
        return ""

    ETX: Annotated[str, "end of text"] = "\x03"
    EOT: Annotated[str, "end of transmission"] = "\x04"
    ENQ: Annotated[str, "enquiry"] = "\x05"
    ACK: Annotated[str, "acknowlege"] = "\x06"

    @property
    def BEL(self: Self) -> str:
        match self.mode:
            case Mode.C0 | Mode.C1:
                return "\x07"
            case Mode.Bash:
                return "\\a"

    BS: Annotated[str, "backspace"] = "\x08"
    HT: Annotated[str, "horizontal tab"] = "\x09"
    LF: Annotated[str, "linefeed"] = "\x0a"
    VT: Annotated[str, "vertical tab"] = "\x0b"
    FF: Annotated[str, "form feed"] = "\x0c"
    CR: Annotated[str, "carriage return"] = "\x0d"
    SO: Annotated[str, "shift out"] = "\x0e"
    SI: Annotated[str, "shift in"] = "\x0f"
    DLE: Annotated[str, "data link escape"] = "\x10"
    DC1: Annotated[str, "device control one"] = "\x11"
    XON = DC1
    DC2: Annotated[str, "device control two"] = "\x12"
    DC3: Annotated[str, "device control three"] = "\x13"
    XOFF = DC3
    DC4: Annotated[str, "device control four"] = "\x14"
    NAK: Annotated[str, "negative acknowledge"] = "\x15"
    SYN: Annotated[str, "synchronous idle"] = "\x16"
    ETB: Annotated[str, "end of transmission block"] = "\x17"
    CAN: Annotated[str, "cancel"] = "\x18"
    EM: Annotated[str, "end of medium"] = "\x19"
    SUB: Annotated[str, "substitute"] = "\x1a"

    @property
    def ESC(self: Self) -> str:
        """escape"""
        match self.mode:
            case Mode.C0 | Mode.C1:
                return "\x1b"
            case Mode.Bash:
                return "\\e"

    FS: Annotated[str, "file separator"] = "\x1c"
    GS: Annotated[str, "group separator"] = "\x1d"
    RS: Annotated[str, "record separator"] = "\x1e"
    US: Annotated[str, "unit separator"] = "\x1f"

    DEL: Annotated[str, "delete"] = "\x7f"

    @property
    def DCS(self: Self) -> str:
        """Device Control String"""
        match self.mode:
            case Mode.C0 | Mode.Bash:
                return f"{self.ESC}P"
            case Mode.C1:
                return "\x8d"

    @property
    def CSI(self) -> str:
        """Control Sequence Introducer"""
        match self.mode:
            case Mode.C0 | Mode.Bash:
                return f"{self.ESC}["
            case Mode.C1:
                return "\x9b"

    @property
    def ST(self) -> str:
        match self.mode:
            case Mode.C0:
                return f"{self.ESC}\\"
            case Mode.C1:
                return "\x9c"
            case Mode.Bash:
                return self.BEL

    @property
    def OSC(self: Self) -> str:
        """Operating System Command"""
        match self.mode:
            case Mode.C0 | Mode.Bash:
                return f"{self.ESC}]"
            case Mode.C1:
                return "\x9d"

    @readline_wrap
    def GraphicsRendition(self: Self, *args: GraphicsRendtion) -> str:
        result = self.CSI
        for index, g in enumerate(args):
            if index > 0:
                result += ";"
            match g:
                case Reset():
                    result += "0"

                case Style():
                    result += "{g}"

                case ResetStyle():
                    result += "{g}"

                case UnderlineStyle():
                    result += "4:{g}"

                case Color():
                    match g.color:
                        case Color8():
                            match g.layer:
                                case Layer.FOREGROUND:
                                    if g.bright:
                                        offset = 90
                                    else:
                                        offset = 30
                                case Layer.BACKGROUND:
                                    if g.bright:
                                        offset = 100
                                    else:
                                        offset = 40
                                case Layer.UNDERLINE:
                                    offset = 50
                            result += f"{g.color + offset}"

                        case Color256():
                            match g.layer:
                                case Layer.FOREGROUND:
                                    code = 38
                                case Layer.BACKGROUND:
                                    code = 48
                                case Layer.UNDERLINE:
                                    code = 58
                            result += f"{code};5;{g.color.color}"

                        case ColorRGB():
                            match g.layer:
                                case Layer.FOREGROUND:
                                    code = 38
                                case Layer.BACKGROUND:
                                    code = 48
                                case Layer.UNDERLINE:
                                    code = 58
                            result += (
                                f"{code};2;{g.color.red};{g.color.green};{g.color.blue}"
                            )

        result += "m"
        return result

    def CursorUp(self: Self, n: int) -> str:
        return f"{self.CSI}{n}A"

    def CursorDown(self: Self, n: int) -> str:
        return f"{self.CSI}{n}B"

    def CursorForward(self: Self, n: int) -> str:
        return f"{self.CSI}{n}C"

    def CursorBack(self: Self, n: int) -> str:
        return f"{self.CSI}{n}D"

    def CursorNextLine(self: Self, n: int) -> str:
        return f"{self.CSI}{n}E"

    def CursorPreviousLine(self: Self, n: int) -> str:
        return f"{self.CSI}{n}F"

    def CursorHorizontalAbsolute(self: Self, n: int, alternate: bool = False) -> str:
        if alternate:
            return f"{self.CSI}{n}f"
        return f"{self.CSI}{n}G"

    def CursorPosition(self: Self, n: int, m: int, alternate: bool = False) -> str:
        if alternate:
            return f"{self.CSI}{n};{m}f"
        return f"{self.CSI}{n};{m}H"

    def EraseInDisplay(self: Self, mode: EraseDisplayMode) -> str:
        return f"{self.CSI}{mode}J"

    def EraseInLine(self: Self, mode: EraseLineMode) -> str:
        return f"{self.CSI}{mode}K"

    def ScrollUp(self: Self, n: int) -> str:
        return f"{self.CSI}{n}S"

    def ScrollDown(self: Self, n: int) -> str:
        return f"{self.CSI}{n}T"

    @property
    def SaveCursorPosition(self: Self) -> str:
        return f"{self.CSI}s"

    @property
    def RestoreCursorPosition(self: Self) -> str:
        return f"{self.CSI}u"

    def Title(
        self: Self, text: str, icon_name: bool = True, window_title: bool = True
    ) -> str:
        code = 0
        match (icon_name, window_title):
            case (False, False):
                raise ValueError("icon_name and window_tile can't both be false")
            case (False, True):
                code = 2
            case (True, False):
                code = 1
            case (True, True):
                code = 0
        return self._wrap(f"{self.OSC}{code};{text}{self.ST}")

    def Hyperlink(self: Self, link: str, text: str) -> str:
        return (
            self._wrap(f"{self.OSC}8;;{link}{self.ST}")
            + text
            + self._wrap(f"{self.OSC}8;;{self.ST}")
        )

    @readline_wrap
    def Notification(self: Self, text: str) -> str:
        """https://iterm2.com/documentation-escape-codes.html"""
        return f"{self.OSC}9;{text}{self.ST}"

    @readline_wrap
    def SetMark(self: Self) -> str:
        """https://iterm2.com/documentation-escape-codes.html"""
        return f"{self.OSC}1337;SetMark{self.ST}"

    @readline_wrap
    def StealFocus(self: Self) -> str:
        """https://iterm2.com/documentation-escape-codes.html"""
        return f"{self.OSC}1337;StealFocus{self.ST}"

    @readline_wrap
    def CurrentDirectory(self: Self, text: str) -> str:
        """
        Set the current directory

        https://iterm2.com/documentation-escape-codes.html
        """
        return f"{self.OSC}1337;CurrentDir={text}{self.ST}"

    @readline_wrap
    def CopyToClipboard(self: Self, text: str) -> str:
        """https://iterm2.com/documentation-escape-codes.html"""
        data = b64encode(text.encode("utf-8")).decode("utf-8")
        return f"{self.OSC}52;Pc;{data}{self.ST}"

    @readline_wrap
    def ChangeProfile(self: Self, profile: str) -> str:
        """https://iterm2.com/documentation-escape-codes.html"""
        return f"{self.OSC}1337;SetProfile={profile}{self.ST}"

    @readline_wrap
    def CursorShape(self: Self, shape: CursorShape) -> str:
        """Set the cursor shape

        https://iterm2.com/documentation-escape-codes.html
        """
        return f"{self.OSC}1337;CursorShape={shape}{self.ST}"

    @readline_wrap
    def ActiveStatusDisplay(self: Self, mode: ActiveStatusDisplayMode) -> str:
        return f"{self.CSI}{mode}$}}"

    @readline_wrap
    def StatusLineType(self: Self, mode: StatusLineType) -> str:
        return f"{self.CSI}{mode}$~"

    @check_annotations
    @readline_wrap
    def DesktopNotification(
        self: Self,
        identifier: Annotated[str, RegexMatch(re.compile(r"\A[a-zA-Z0-9-_+.]\z"))],
        title: str,
        body: str,
    ) -> str:
        def format(
            part: DesktopNotificationPart, text: str, done: DesktopNotificationDone
        ) -> str:
            return f"{self.OSC}99;i={identifier};e=0;p={part};d={done};{text}{self.ST}"

        result = ""
        stride = 256

        for part in [title[i : i + stride] for i in range(0, len(title), stride)]:
            done = DesktopNotificationDone.NOT_DONE
            if len(part) < stride and len(body) == 0:
                done = DesktopNotificationDone.DONE

            result += format(DesktopNotificationPart.TITLE, part, done)

        for part in [body[i : i + stride] for i in range(0, len(body), stride)]:
            done = DesktopNotificationDone.NOT_DONE
            if len(part) < stride:
                done = DesktopNotificationDone.DONE

            result += format(DesktopNotificationPart.TITLE, part, done)

        return result