colorisedstring.py

from enum import Enum
from typing import Optional, TypeVar, Union

T = TypeVar("T")

CSI = "\x1b["


class AnsiCodes(Enum):
    def __str__(self) -> str:
        return str(self.value)


class AnsiFore(AnsiCodes):
    BLACK = 30
    RED = 31
    GREEN = 32
    YELLOW = 33
    BLUE = 34
    MAGENTA = 35
    CYAN = 36
    WHITE = 37
    RESET = 39
    LIGHTBLACK_EX = 90
    LIGHTRED_EX = 91
    LIGHTGREEN_EX = 92
    LIGHTYELLOW_EX = 93
    LIGHTBLUE_EX = 94
    LIGHTMAGENTA_EX = 95
    LIGHTCYAN_EX = 96
    LIGHTWHITE_EX = 97


class AnsiBack(AnsiCodes):
    BLACK = 40
    RED = 41
    GREEN = 42
    YELLOW = 43
    BLUE = 44
    MAGENTA = 45
    CYAN = 46
    WHITE = 47
    RESET = 49
    LIGHTBLACK_EX = 100
    LIGHTRED_EX = 101
    LIGHTGREEN_EX = 102
    LIGHTYELLOW_EX = 103
    LIGHTBLUE_EX = 104
    LIGHTMAGENTA_EX = 105
    LIGHTCYAN_EX = 106
    LIGHTWHITE_EX = 107


class AnsiStyle(AnsiCodes):
    BRIGHT = 1
    DIM = 2
    NORMAL = 22
    RESET_ALL = 0


class ColorizedString(str):
    def colorize(
        self,
        fore: Optional[AnsiFore] = None,
        back: Optional[AnsiBack] = None,
        style: Optional[AnsiStyle] = None,
    ) -> "ColorizedString":
        self._fore = fore
        self._back = back
        self._style = style
        return self

    def __str__(self) -> str:
        value = super().__str__()
        fore = getattr(self, "_fore", None)
        back = getattr(self, "_back", None)
        style = getattr(self, "_style", None)
        if fore is not None:
            value = f"\x1b[1m{CSI}{fore}m{value}"
        if back is not None:
            value = f"{CSI}{back}m{value}"
        if style is not None:
            value = f"{CSI}{style}m{value}"
        return value + "\x1b[0m"

    def __repr__(self) -> str:
        return repr(self.__str__())

    def __getitem__(self, i: Union[int, slice]) -> "ColorizedString":
        return ColorizedString(super().__getitem__(i)).colorize(
            getattr(self, "_fore", None),
            getattr(self, "_back", None),
            getattr(self, "_style", None),
        )


if __name__ == "__main__":
    hello = ColorizedString("Hello").colorize(
        AnsiFore.WHITE, AnsiBack.MAGENTA, AnsiStyle.BRIGHT
    )
    world = ColorizedString("world!").colorize(
        AnsiFore.LIGHTBLACK_EX, AnsiBack.LIGHTGREEN_EX, AnsiStyle.DIM
    )
    assert hello == "Hello"
    assert world == "world!"
    print(f"{hello} {world}")