message pills

pills.py

import dataclasses
import io
import os
import re
import shutil
import textwrap
import typing


Color8Type: typing.TypeAlias = typing.Literal[0, 1, 2, 3, 4, 5, 6, 7, 9]


@dataclasses.dataclass(slots=True)
class Box:
    content: str
    border_color: Color8Type = dataclasses.field(default=9, kw_only=True)
    margin_top: int = dataclasses.field(default=0, kw_only=True)
    margin_bottom: int = dataclasses.field(default=0, kw_only=True)
    margin_left: int = dataclasses.field(default=1, kw_only=True)
    margin_right: int = dataclasses.field(default=1, kw_only=True)
    padding_top: int = dataclasses.field(default=0, kw_only=True)
    padding_bottom: int = dataclasses.field(default=0, kw_only=True)
    padding_left: int = dataclasses.field(default=1, kw_only=True)
    padding_right: int = dataclasses.field(default=1, kw_only=True)

    top_left: typing.ClassVar[str] = "╭"
    top_right: typing.ClassVar[str] = "╮"
    bottom_left: typing.ClassVar[str] = "╰"
    bottom_right: typing.ClassVar[str] = "╯"
    horizontal: typing.ClassVar[str] = "─"
    vertical: typing.ClassVar[str] = "│"

    def calc_widths(self, term_width: int | None = None) -> tuple[int, int]:
        _term_width = (
            shutil.get_terminal_size().columns if term_width is None else term_width
        )

        padding_box_width = _term_width - self.margin_left - self.margin_right - 2
        content_box_width = padding_box_width - self.padding_left - self.padding_right

        return padding_box_width, content_box_width

    def get_colored_border(self, border: str) -> str:
        return f"\x1b[3{self.border_color}m{border}\x1b[39m"

    def get_contents(self) -> str:
        return self.content

    def write_line(
        self,
        buffer: typing.TextIO,
        content: str,
        border_left: str,
        border_right: str,
    ) -> None:
        buffer.write(" " * self.margin_left)
        buffer.write(self.get_colored_border(border_left))
        buffer.write(content)
        buffer.write(self.get_colored_border(border_right))
        buffer.write(" " * self.margin_right)
        buffer.write("\n")

    def render(self) -> str:
        buffer = io.StringIO()
        term_width, _ = shutil.get_terminal_size()

        padding_box_width, content_box_width = self.calc_widths(term_width)

        buffer.write("\n" * self.margin_top)

        self.write_line(
            buffer,
            self.get_colored_border(self.horizontal * padding_box_width),
            self.top_left,
            self.top_right,
        )

        for _ in range(self.padding_top):
            self.write_line(
                buffer,
                " " * padding_box_width,
                self.vertical,
                self.vertical,
            )

        for line in textwrap.wrap(self.get_contents(), content_box_width):
            trailing = " " * (content_box_width - true_length(line))
            content = (
                " " * self.padding_left + line + trailing + " " * self.padding_right
            )
            self.write_line(buffer, content, self.vertical, self.vertical)

        for _ in range(self.padding_bottom):
            self.write_line(
                buffer,
                " " * padding_box_width,
                self.vertical,
                self.vertical,
            )

        self.write_line(
            buffer,
            self.get_colored_border(self.horizontal * padding_box_width),
            self.bottom_left,
            self.bottom_right,
        )

        buffer.write("\n" * self.margin_bottom)

        return buffer.getvalue()


@dataclasses.dataclass(slots=True)
class WarnBox(Box):
    border_color: Color8Type = dataclasses.field(init=False, default=3)

    synopsis: typing.ClassVar[str] = "Warning"

    def get_contents(self) -> str:
        synopsis = f"\x1b[1;43m {self.synopsis} \x1b[22;49m "

        return synopsis + self.content


@dataclasses.dataclass(slots=True)
class ErrorBox(Box):
    border_color: Color8Type = dataclasses.field(init=False, default=1)

    synopsis: typing.ClassVar[str] = "Error"

    def get_contents(self) -> str:
        synopsis = f"\x1b[1;41m {self.synopsis} \x1b[22;49m "

        return synopsis + self.content


def strip_escape_sequences(string: str) -> str:
    return re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]").sub("", string)


def true_length(string: str) -> int:
    return len(strip_escape_sequences(string))


def main() -> int:
    print(
        WarnBox(
            "This is deprecated and will be removed in the future.",
            margin_top=1,
        ).render(),
        end="",
    )
    print(ErrorBox("Actually this has already been removed lol").render(), end="")

    return os.EX_OK


if __name__ == "__main__":
    raise SystemExit(main())