X-Git-Url: https://wannabe.guru.org/gitweb/?a=blobdiff_plain;f=ansi.py;h=c13e3522fe187efad36ae73de6d174ee76c4b2aa;hb=31c81f6539969a5eba864d3305f9fb7bf716a367;hp=9e31b811ab978fa1ae81c3974ba991700287d867;hpb=e6f32fdd9b373dfcd100c7accb41f57d83c2f0a1;p=python_utils.git

diff --git a/ansi.py b/ansi.py
index 9e31b81..c13e352 100755
--- a/ansi.py
+++ b/ansi.py
@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 
-from abc import abstractmethod
 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
@@ -1890,19 +1890,19 @@ def bg(
 
 class StdoutInterceptor(io.TextIOBase):
     def __init__(self):
-        self.saved_stdout: Optional[io.TextIOBase] = None
+        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