from __future__ import annotations from .prettytable import PrettyTable try: from colorama import init init() except ImportError: pass RESET_CODE = "\x1b[0m" class Theme: def __init__( self, default_color: str = "", vertical_char: str = "|", vertical_color: str = "", horizontal_char: str = "-", horizontal_color: str = "", junction_char: str = "+", junction_color: str = "", ) -> None: self.default_color = Theme.format_code(default_color) self.vertical_char = vertical_char self.vertical_color = Theme.format_code(vertical_color) self.horizontal_char = horizontal_char self.horizontal_color = Theme.format_code(horizontal_color) self.junction_char = junction_char self.junction_color = Theme.format_code(junction_color) @staticmethod def format_code(s: str) -> str: """Takes string and intelligently puts it into an ANSI escape sequence""" if s.strip() == "": return "" elif s.startswith("\x1b["): return s else: return f"\x1b[{s}m" class Themes: DEFAULT = Theme() OCEAN = Theme( default_color="96", vertical_color="34", horizontal_color="34", junction_color="36", ) class ColorTable(PrettyTable): def __init__(self, field_names=None, **kwargs) -> None: super().__init__(field_names=field_names, **kwargs) # TODO: Validate option self.theme = kwargs.get("theme") or Themes.DEFAULT @property def theme(self) -> Theme: return self._theme @theme.setter def theme(self, value: Theme): self._theme = value self.update_theme() def update_theme(self) -> None: theme = self._theme self._vertical_char = ( theme.vertical_color + theme.vertical_char + RESET_CODE + theme.default_color ) self._horizontal_char = ( theme.horizontal_color + theme.horizontal_char + RESET_CODE + theme.default_color ) self._junction_char = ( theme.junction_color + theme.junction_char + RESET_CODE + theme.default_color ) def get_string(self, **kwargs) -> str: return super().get_string(**kwargs) + RESET_CODE