Don't be tricked by speaker groups living on the same IP address

[python_utils.git] / ansi.py

diff --git a/ansi.py b/ansi.py
index 476f0a4dd9f80c111567a6b8c5395ac17f2686a2..950a0c4a3c809587bb7c8324a884c62d83b553ab 100755 (executable)
--- 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__)
 
@@ -1620,7 +1625,7 @@ def italic() -> str:
 
 
 def italics() -> str:
-    return "\e[3m"
+    return italic()
 
 
 def underline() -> str:
@@ -1632,7 +1637,7 @@ def strikethrough() -> str:
 
 
 def strike_through() -> str:
-    return "\e[9m"
+    return strikethrough()
 
 
 def is_16color(num: int) -> bool:
@@ -1725,6 +1730,7 @@ def _find_color_by_name(name: str) -> Tuple[int, int, int]:
     return rgb
 
 
+@logging_utils.squelch_repeated_log_messages(1)
 def fg(name: Optional[str] = "",
        red: Optional[int] = None,
        green: Optional[int] = None,
@@ -1732,6 +1738,22 @@ def fg(name: Optional[str] = "",
        *,
        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(
@@ -1777,6 +1799,19 @@ 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:
@@ -1812,6 +1847,19 @@ def bg(name: Optional[str] = "",
        *,
        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(
@@ -1842,18 +1890,54 @@ 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()