X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=ansi.py;h=a8b84fc0c2dd0acce18e96d21d86eb2b2da409b7;hb=3ca9b4d16433af8da5d2de7f4a2338b56b5428d5;hp=4c580c0a0262ad133a71f6def6ca63a7d003dd8f;hpb=11eeb8574b7b4620ac6fd440cb251f8aa2458f5b;p=python_utils.git diff --git a/ansi.py b/ansi.py index 4c580c0..a8b84fc 100755 --- a/ansi.py +++ b/ansi.py @@ -1,9 +1,25 @@ #!/usr/bin/env python3 +# © Copyright 2021-2022, Scott Gasch + +"""A bunch of color names mapped into RGB tuples and some methods for +setting the text color, background, etc... using ANSI escape +sequences. +""" + +import contextlib import difflib +import io import logging +import re import sys -from typing import Dict, Optional, Tuple +from abc import abstractmethod +from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple + +from overrides import overrides + +import logging_utils +import string_utils logger = logging.getLogger(__name__) @@ -1594,54 +1610,67 @@ COLOR_NAMES_TO_RGB: Dict[str, Tuple[int, int, int]] = { def clear() -> str: + """Clear screen ANSI escape sequence""" return "" def clear_screen() -> str: + """Clear screen ANSI escape sequence""" return "" def reset() -> str: + """Reset text attributes to 'normal'""" return "" def normal() -> str: + """Reset text attributes to 'normal'""" return "" def bold() -> str: + """Set text to bold""" return "" def italic() -> str: + """Set text to italic""" return "" def italics() -> str: + """Set text to italic""" return italic() def underline() -> str: + """Set text to underline""" return "" def strikethrough() -> str: + """Set text to strikethrough""" return "" def strike_through() -> str: + """Set text to strikethrough""" return strikethrough() def is_16color(num: int) -> bool: - return num == 255 or num == 128 + """Is num a valid 16 color number?""" + return num in (255, 128) def is_216color(num: int) -> bool: + """Is num a valid 256 color number?""" return num in set([0, 95, 135, 175, 223, 255]) def _simple_color_number(red: int, green: int, blue: int) -> int: + """Construct a simple color number""" r = red > 0 g = green > 0 b = blue > 0 @@ -1649,6 +1678,7 @@ def _simple_color_number(red: int, green: int, blue: int) -> int: def fg_16color(red: int, green: int, blue: int) -> str: + """Set foreground color using 16 color mode""" code = _simple_color_number(red, green, blue) + 30 bright_count = 0 if red > 128: @@ -1663,6 +1693,7 @@ def fg_16color(red: int, green: int, blue: int) -> str: def bg_16color(red: int, green: int, blue: int) -> str: + """Set background using 16 color mode""" code = _simple_color_number(red, green, blue) + 40 bright_count = 0 if red > 128: @@ -1691,6 +1722,7 @@ def _pixel_to_216color(n: int) -> int: def fg_216color(red: int, green: int, blue: int) -> str: + """Set foreground using 216 color mode""" r = _pixel_to_216color(red) g = _pixel_to_216color(green) b = _pixel_to_216color(blue) @@ -1699,6 +1731,7 @@ def fg_216color(red: int, green: int, blue: int) -> str: def bg_216color(red: int, green: int, blue: int) -> str: + """Set background using 216 color mode""" r = _pixel_to_216color(red) g = _pixel_to_216color(green) b = _pixel_to_216color(blue) @@ -1707,10 +1740,12 @@ def bg_216color(red: int, green: int, blue: int) -> str: def fg_24bit(red: int, green: int, blue: int) -> str: + """Set foreground using 24bit color mode""" return f"[38;2;{red};{green};{blue}m" def bg_24bit(red: int, green: int, blue: int) -> str: + """Set background using 24bit color mode""" return f"[48;2;{red};{green};{blue}m" @@ -1723,14 +1758,39 @@ 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: - import string_utils +@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. + + Args: + name: the name of the color to set + red: the color to set's red component value + green: the color to set's green component value + blue: the color to set's blue component value + force_16color: force fg to use 16 color mode + force_216color: force fg to use 216 color mode + + Returns: + String containing the ANSI escape sequence to set desired foreground + + >>> import string_utils as su + >>> su.to_base64(fg('blue')) + b'G1szODs1OzIxbQ==\\n' + """ + 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) @@ -1740,7 +1800,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: @@ -1749,14 +1809,10 @@ def fg(name: Optional[str] = "", green = 0 if blue is None: blue = 0 - if ( - is_16color(red) and is_16color(green) and is_16color(blue) - ) or force_16color: + if (is_16color(red) and is_16color(green) and is_16color(blue)) or force_16color: logger.debug("Using 16-color strategy") return fg_16color(red, green, blue) - if ( - is_216color(red) and is_216color(green) and is_216color(blue) - ) or force_216color: + if (is_216color(red) and is_216color(green) and is_216color(blue)) or force_216color: logger.debug("Using 216-color strategy") return fg_216color(red, green, blue) logger.debug("Using 24-bit color strategy") @@ -1769,16 +1825,36 @@ 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]: - import string_utils +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. + + Args: + name: the name of the color to contrast + red: the color to contrast's red component value + green: the color to contrast's green component value + blue: the color to contrast's blue component value + + Returns: + An RGB tuple containing a contrasting color + >>> pick_contrasting_color(None, 20, 20, 20) + (255, 255, 255) + + >>> pick_contrasting_color("white") + (0, 0, 0) + + """ if name is not None and string_utils.is_full_string(name): rgb = _find_color_by_name(name) else: @@ -1791,30 +1867,46 @@ def pick_contrasting_color(name: Optional[str] = "", def guess_name(name: str) -> str: + """Try to guess what color the user is talking about""" 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 assert best_guess is not None - logger.debug(f"Best guess at color name is {best_guess}") + logger.debug("Best guess at color name is %s", best_guess) 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: - import string_utils +@logging_utils.squelch_repeated_log_messages(1) +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. + + Args: + name: the name of the color to set + red: the color to set's red component value + green: the color to set's green component value + blue: the color to set's blue component value + force_16color: force bg to use 16 color mode + force_216color: force bg to use 216 color mode + + >>> import string_utils as su + >>> su.to_base64(bg("red")) # b'\x1b[48;5;196m' + b'G1s0ODs1OzE5Nm0=\\n' + """ + 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) @@ -1824,7 +1916,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 @@ -1832,22 +1924,66 @@ def bg(name: Optional[str] = "", green = 0 if blue is None: blue = 0 - if ( - is_16color(red) and is_16color(green) and is_16color(blue) - ) or force_16color: + if (is_16color(red) and is_16color(green) and is_16color(blue)) or force_16color: logger.debug("Using 16-color strategy") return bg_16color(red, green, blue) - if ( - is_216color(red) and is_216color(green) and is_216color(blue) - ) or force_216color: + if (is_216color(red) and is_216color(green) and is_216color(blue)) or force_216color: logger.debug("Using 216-color strategy") return bg_216color(red, green, blue) logger.debug("Using 24-bit color strategy") return bg_24bit(red, green, blue) +class StdoutInterceptor(io.TextIOBase, contextlib.AbstractContextManager): + """An interceptor for data written to stdout. Use as a context.""" + + def __init__(self): + super().__init__() + self.saved_stdout: io.TextIO = None + self.buf = '' + + @abstractmethod + def write(self, s: str): + pass + + def __enter__(self): + self.saved_stdout = sys.stdout + sys.stdout = self + return self + + def __exit__(self, *args) -> Literal[False]: + sys.stdout = self.saved_stdout + print(self.buf) + return False + + +class ProgrammableColorizer(StdoutInterceptor): + """A colorizing interceptor; pass it re.Patterns -> methods that do + something (usually add color to) the match. + + """ + + def __init__( + self, + patterns: Iterable[Tuple[re.Pattern, Callable[[Any, re.Pattern], str]]], + ): + super().__init__() + self.patterns = list(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: @@ -1856,6 +1992,6 @@ if __name__ == '__main__': _ = 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()}') + print(f'{f}{xb}{possibility}{reset()}\t\t\t' f'{b}{xf}{possibility}{reset()}') + main()