More cleanup, yey!
[python_utils.git] / ansi.py
diff --git a/ansi.py b/ansi.py
index d30ae27e625dbe585d82d3392837b39bf436f888..03f8fd27473c4e295cdafbb9e7122b4cec296453 100755 (executable)
--- a/ansi.py
+++ b/ansi.py
@@ -1,16 +1,22 @@
 #!/usr/bin/env python3
 
-from abc import abstractmethod
+"""A bunch of color names mapped into RGB tuples and some methods for
+setting the text color, background, etc... using ANSI escape
+sequences."""
+
 import difflib
 import io
 import logging
 import re
 import sys
+from abc import abstractmethod
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
 from overrides import overrides
 
 import logging_utils
+import string_utils
+
 
 logger = logging.getLogger(__name__)
 
@@ -1641,7 +1647,7 @@ def strike_through() -> str:
 
 
 def is_16color(num: int) -> bool:
-    return num == 255 or num == 128
+    return num in (255, 128)
 
 
 def is_216color(num: int) -> bool:
@@ -1731,15 +1737,26 @@ def _find_color_by_name(name: str) -> Tuple[int, int, int]:
 
 
 @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:
-    import string_utils
-
+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'
+
+    """
     if name is not None and name == 'reset':
         return '\033[39m'
 
@@ -1751,7 +1768,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:
@@ -1760,14 +1777,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")
@@ -1780,16 +1793,27 @@ 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.
+
+    >>> 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:
@@ -1805,28 +1829,33 @@ 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
     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.
+
+    >>> 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'
 
@@ -1838,7 +1867,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
@@ -1846,14 +1875,10 @@ 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")
@@ -1861,29 +1886,40 @@ def bg(name: Optional[str] = "",
 
 
 class StdoutInterceptor(io.TextIOBase):
+    """An interceptor for data written to stdout.  Use as a context.
+
+    """
     def __init__(self):
-        self.saved_stdout: Optional[io.TextIOBase] = None
+        super().__init__()
+        self.saved_stdout: io.TextIO = None
         self.buf = ''
 
     @abstractmethod
     def write(self, s: str):
         pass
 
-    def __enter__(self) -> None:
+    def __enter__(self):
         self.saved_stdout = sys.stdout
         sys.stdout = self
-        return None
+        return self
 
-    def __exit__(self, *args) -> bool:
+    def __exit__(self, *args) -> Optional[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]]]):
+    """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 = [_ for _ in patterns]
+        self.patterns = list(patterns)
 
     @overrides
     def write(self, s: str):
@@ -1894,6 +1930,10 @@ class ProgrammableColorizer(StdoutInterceptor):
 
 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:
@@ -1902,6 +1942,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()