X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=ansi.py;h=dee52424fefa6bd80daccc3d8830c0a707e12b6a;hb=f995224a77488f3469c16a3f504dcf7a1393d834;hp=03f8fd27473c4e295cdafbb9e7122b4cec296453;hpb=5c212d7639f62fcb936f9d7a0bbe704a9f7b213d;p=python_utils.git

diff --git a/ansi.py b/ansi.py
index 03f8fd2..dee5242 100755
--- a/ansi.py
+++ b/ansi.py
@@ -1,23 +1,27 @@
 #!/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."""
+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__)
 
 # https://en.wikipedia.org/wiki/ANSI_escape_code
@@ -1885,10 +1889,9 @@ def bg(
     return bg_24bit(red, green, blue)
 
 
-class StdoutInterceptor(io.TextIOBase):
-    """An interceptor for data written to stdout.  Use as a context.
+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
@@ -1903,10 +1906,10 @@ 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):
@@ -1914,6 +1917,7 @@ class ProgrammableColorizer(StdoutInterceptor):
     something (usually add color to) the match.
 
     """
+
     def __init__(
         self,
         patterns: Iterable[Tuple[re.Pattern, Callable[[Any, re.Pattern], str]]],
@@ -1929,6 +1933,7 @@ class ProgrammableColorizer(StdoutInterceptor):
 
 
 if __name__ == '__main__':
+
     def main() -> None:
         import doctest