From 0d30debeac32a04383d95af18ed334b080de148f Mon Sep 17 00:00:00 2001 From: Sijmen Schoon Date: Wed, 4 Sep 2019 16:39:36 +0200 Subject: [PATCH] Improve color.py --- color.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/color.py b/color.py index 8bdb0fd..432782f 100644 --- a/color.py +++ b/color.py @@ -20,26 +20,34 @@ class Color: BRIGHT_CYAN = 14 BRIGHT_WHITE = 15 - def __init__(self, foreground: Optional[int], - background: Optional[int] = None) -> None: + def __init__( + self, foreground: Optional[int], background: Optional[int] = None + ) -> None: self.foreground = foreground self.background = background - def print(self, *args: Any, end: str = '\n') -> None: + def print(self, *args: Any, end: str = "\n") -> None: print(self.format(*args), end=end) def format(self, *args: Any) -> str: - colors = ['0'] - + # Create the color string. + colors = [] if self.foreground is not None: colors.append(f"38;5;{self.foreground}") - if self.background is not None: colors.append(f"48;5;{self.background}") + color_str = ";".join(colors) - color_str = ';'.join(colors) - text = ' '.join(str(a) for a in args) + text = " ".join(str(a) for a in args) if not text: - return '' + return "" - return f'\033[{color_str}m{text}\033[0m' + # Reset the color as necessary. + reset = [] + if self.foreground is not None: + reset.append("39") + if self.background is not None: + reset.append("49") + reset_str = ";".join(reset) + + return f"\033[{color_str}m{text}\033[{reset_str}m"