X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=ansi.py;h=950a0c4a3c809587bb7c8324a884c62d83b553ab;hb=8d6abef443543013d7b37cda1375a28f040d24a6;hp=dfbf48a0b0858c50cb2039ec14d169115a86599b;hpb=d92df9c03f8510042ee1f8949d0786761f01d35e;p=python_utils.git

diff --git a/ansi.py b/ansi.py
index dfbf48a..950a0c4 100755
--- a/ansi.py
+++ b/ansi.py
@@ -8,6 +8,8 @@ import re
 import sys
 from typing import Any, Callable, Dict, Iterable, Optional, Tuple
 
+from overrides import overrides
+
 import logging_utils
 
 logger = logging.getLogger(__name__)
@@ -1736,8 +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(
@@ -1783,6 +1799,17 @@ 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):
@@ -1820,8 +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(
@@ -1858,7 +1896,7 @@ class StdoutInterceptor(io.TextIOBase):
         self.buf = ''
 
     @abstractmethod
-    def write(self, s):
+    def write(self, s: str):
         pass
 
     def __enter__(self) -> None:
@@ -1873,10 +1911,14 @@ class StdoutInterceptor(io.TextIOBase):
 
 
 class ProgrammableColorizer(StdoutInterceptor):
-    def __init__(self, patterns: Iterable[Tuple[re.Pattern, Callable[[Any, re.Pattern], str]]]):
+    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)
@@ -1885,6 +1927,9 @@ 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: