Improve docs in ANSI.py
[python_utils.git] / ansi.py
diff --git a/ansi.py b/ansi.py
index 1633fddbcb31714d3ae7342daca37c1c4034b4c6..a8b84fc0c2dd0acce18e96d21d86eb2b2da409b7 100755 (executable)
--- a/ansi.py
+++ b/ansi.py
@@ -1,16 +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 abc import abstractmethod
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, Literal, Optional, Tuple
 
 from overrides import overrides
 
 import logging_utils
+import string_utils
 
 logger = logging.getLogger(__name__)
 
@@ -1601,54 +1610,67 @@ COLOR_NAMES_TO_RGB: Dict[str, Tuple[int, int, int]] = {
 
 
 def clear() -> str:
+    """Clear screen ANSI escape sequence"""
     return "\e[H\e[2J"
 
 
 def clear_screen() -> str:
+    """Clear screen ANSI escape sequence"""
     return "\e[H\e[2J"
 
 
 def reset() -> str:
+    """Reset text attributes to 'normal'"""
     return "\e[m"
 
 
 def normal() -> str:
+    """Reset text attributes to 'normal'"""
     return "\e[m"
 
 
 def bold() -> str:
+    """Set text to bold"""
     return "\e[1m"
 
 
 def italic() -> str:
+    """Set text to italic"""
     return "\e[3m"
 
 
 def italics() -> str:
+    """Set text to italic"""
     return italic()
 
 
 def underline() -> str:
+    """Set text to underline"""
     return "\e[4m"
 
 
 def strikethrough() -> str:
+    """Set text to strikethrough"""
     return "\e[9m"
 
 
 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
@@ -1656,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:
@@ -1670,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:
@@ -1698,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)
@@ -1706,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)
@@ -1714,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"\e[38;2;{red};{green};{blue}m"
 
 
 def bg_24bit(red: int, green: int, blue: int) -> str:
+    """Set background using 24bit color mode"""
     return f"\e[48;2;{red};{green};{blue}m"
 
 
@@ -1746,13 +1774,21 @@ def fg(
     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'
-
     """
-    import string_utils
-
     if name is not None and name == 'reset':
         return '\033[39m'
 
@@ -1803,6 +1839,15 @@ def pick_contrasting_color(
     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)
 
@@ -1810,8 +1855,6 @@ def pick_contrasting_color(
     (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:
@@ -1824,6 +1867,7 @@ def pick_contrasting_color(
 
 
 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:
@@ -1832,7 +1876,7 @@ def guess_name(name: str) -> str:
             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
 
 
@@ -1849,13 +1893,18 @@ def bg(
     """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'
-
     """
-    import string_utils
-
     if name is not None and name == 'reset':
         return '\033[49m'
 
@@ -1885,8 +1934,11 @@ def bg(
     return bg_24bit(red, green, blue)
 
 
-class StdoutInterceptor(io.TextIOBase):
+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 = ''
 
@@ -1899,19 +1951,24 @@ class StdoutInterceptor(io.TextIOBase):
         sys.stdout = self
         return self
 
-    def __exit__(self, *args) -> Optional[bool]:
+    def __exit__(self, *args) -> Literal[False]:
         sys.stdout = self.saved_stdout
         print(self.buf)
-        return None
+        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 = [_ for _ in patterns]
+        self.patterns = list(patterns)
 
     @overrides
     def write(self, s: str):