X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=ansi.py;h=5fde4af56c9ef18254cd562ca19b10d853597e93;hb=36fea7f15ed17150691b5b3ead75450e575229ef;hp=dc9a31542f4551a73d9ab1de341b7fce11e85f47;hpb=dab5654d392f69fb00bed49cf8ffb80f37642ea5;p=python_utils.git diff --git a/ansi.py b/ansi.py index dc9a315..5fde4af 100755 --- a/ansi.py +++ b/ansi.py @@ -1,11 +1,16 @@ #!/usr/bin/env python3 +from abc import abstractmethod import difflib +import io import logging +import re import sys -from typing import Dict, Optional, Tuple +from typing import Any, Callable, Dict, Iterable, Optional, Tuple -import string_utils +from overrides import overrides + +import logging_utils logger = logging.getLogger(__name__) @@ -1725,13 +1730,32 @@ def _find_color_by_name(name: str) -> Tuple[int, int, int]: return rgb -def fg(name: Optional[str] = "", - red: Optional[int] = None, - green: Optional[int] = None, - blue: Optional[int] = None, - *, - force_16color: bool = False, - force_216color: bool = False) -> str: +@logging_utils.squelch_repeated_log_messages(1) +def fg( + name: Optional[str] = "", + red: Optional[int] = None, + green: Optional[int] = None, + blue: Optional[int] = None, + *, + force_16color: bool = False, + force_216color: bool = False, +) -> str: + """Return the ANSI escape sequence to change the foreground color + being printed. Target colors may be indicated by name or R/G/B. + Result will use the 16 or 216 color scheme if force_16color or + force_216color are passed (respectively). Otherwise the code will + do what it thinks best. + + >>> import string_utils as su + >>> su.to_base64(fg('blue')) + b'G1szODs1OzIxbQ==\\n' + + """ + import string_utils + + if name is not None and name == 'reset': + return '\033[39m' + if name is not None and string_utils.is_full_string(name): rgb = _find_color_by_name(name) return fg( @@ -1740,7 +1764,7 @@ def fg(name: Optional[str] = "", rgb[1], rgb[2], force_16color=force_16color, - force_216color=force_216color + force_216color=force_216color, ) if red is None: @@ -1769,14 +1793,29 @@ def _rgb_to_yiq(rgb: Tuple[int, int, int]) -> int: def _contrast(rgb: Tuple[int, int, int]) -> Tuple[int, int, int]: if _rgb_to_yiq(rgb) < 128: - return (0xff, 0xff, 0xff) + return (0xFF, 0xFF, 0xFF) return (0, 0, 0) -def pick_contrasting_color(name: Optional[str] = "", - red: Optional[int] = None, - green: Optional[int] = None, - blue: Optional[int] = None) -> Tuple[int, int, int]: +def pick_contrasting_color( + name: Optional[str] = "", + red: Optional[int] = None, + green: Optional[int] = None, + blue: Optional[int] = None, +) -> Tuple[int, int, int]: + """This method will return a red, green, blue tuple representing a + contrasting color given the red, green, blue of a background + color or a color name of the background color. + + >>> pick_contrasting_color(None, 20, 20, 20) + (255, 255, 255) + + >>> pick_contrasting_color("white") + (0, 0, 0) + + """ + import string_utils + if name is not None and string_utils.is_full_string(name): rgb = _find_color_by_name(name) else: @@ -1792,11 +1831,7 @@ def guess_name(name: str) -> str: best_guess = None max_ratio = None for possibility in COLOR_NAMES_TO_RGB: - r = difflib.SequenceMatcher( - None, - name, - possibility - ).ratio() + r = difflib.SequenceMatcher(None, name, possibility).ratio() if max_ratio is None or r > max_ratio: max_ratio = r best_guess = possibility @@ -1805,13 +1840,28 @@ def guess_name(name: str) -> str: return best_guess -def bg(name: Optional[str] = "", - red: Optional[int] = None, - green: Optional[int] = None, - blue: Optional[int] = None, - *, - force_16color: bool = False, - force_216color: bool = False) -> str: +def bg( + name: Optional[str] = "", + red: Optional[int] = None, + green: Optional[int] = None, + blue: Optional[int] = None, + *, + force_16color: bool = False, + force_216color: bool = False, +) -> str: + """Returns an ANSI color code for changing the current background + color. + + >>> import string_utils as su + >>> su.to_base64(bg("red")) # b'\x1b[48;5;196m' + b'G1s0ODs1OzE5Nm0=\\n' + + """ + import string_utils + + if name is not None and name == 'reset': + return '\033[49m' + if name is not None and string_utils.is_full_string(name): rgb = _find_color_by_name(name) return bg( @@ -1820,7 +1870,7 @@ def bg(name: Optional[str] = "", rgb[1], rgb[2], force_16color=force_16color, - force_216color=force_216color + force_216color=force_216color, ) if red is None: red = 0 @@ -1842,18 +1892,59 @@ def bg(name: Optional[str] = "", return bg_24bit(red, green, blue) -def main() -> None: - name = " ".join(sys.argv[1:]) - for possibility in COLOR_NAMES_TO_RGB: - if name in possibility: - f = fg(possibility) - b = bg(possibility) - _ = pick_contrasting_color(possibility) - xf = fg(None, _[0], _[1], _[2]) - xb = bg(None, _[0], _[1], _[2]) - print(f'{f}{xb}{possibility}{reset()}\t\t\t' - f'{b}{xf}{possibility}{reset()}') +class StdoutInterceptor(io.TextIOBase): + def __init__(self): + self.saved_stdout: Optional[io.TextIOBase] = None + self.buf = '' + + @abstractmethod + def write(self, s: str): + pass + + def __enter__(self) -> None: + self.saved_stdout = sys.stdout + sys.stdout = self + return None + + def __exit__(self, *args) -> bool: + sys.stdout = self.saved_stdout + print(self.buf) + return None + + +class ProgrammableColorizer(StdoutInterceptor): + def __init__( + self, + patterns: Iterable[Tuple[re.Pattern, Callable[[Any, re.Pattern], str]]], + ): + super().__init__() + self.patterns = [_ for _ in patterns] + + @overrides + def write(self, s: str): + for pattern in self.patterns: + s = pattern[0].sub(pattern[1], s) + self.buf += s if __name__ == '__main__': + + def main() -> None: + import doctest + + doctest.testmod() + + name = " ".join(sys.argv[1:]) + for possibility in COLOR_NAMES_TO_RGB: + if name in possibility: + f = fg(possibility) + b = bg(possibility) + _ = pick_contrasting_color(possibility) + xf = fg(None, _[0], _[1], _[2]) + xb = bg(None, _[0], _[1], _[2]) + print( + f'{f}{xb}{possibility}{reset()}\t\t\t' + f'{b}{xf}{possibility}{reset()}' + ) + main()