| # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. |
| import atexit |
| import contextlib |
| import sys |
| |
| from .ansitowin32 import AnsiToWin32 |
| |
| |
| orig_stdout = None |
| orig_stderr = None |
| |
| wrapped_stdout = None |
| wrapped_stderr = None |
| |
| atexit_done = False |
| |
| |
| def reset_all(): |
| if AnsiToWin32 is not None: # Issue #74: objects might become None at exit |
| AnsiToWin32(orig_stdout).reset_all() |
| |
| |
| def init(autoreset=False, convert=None, strip=None, wrap=True): |
| |
| if not wrap and any([autoreset, convert, strip]): |
| raise ValueError('wrap=False conflicts with any other arg=True') |
| |
| global wrapped_stdout, wrapped_stderr |
| global orig_stdout, orig_stderr |
| |
| orig_stdout = sys.stdout |
| orig_stderr = sys.stderr |
| |
| if sys.stdout is None: |
| wrapped_stdout = None |
| else: |
| sys.stdout = wrapped_stdout = \ |
| wrap_stream(orig_stdout, convert, strip, autoreset, wrap) |
| if sys.stderr is None: |
| wrapped_stderr = None |
| else: |
| sys.stderr = wrapped_stderr = \ |
| wrap_stream(orig_stderr, convert, strip, autoreset, wrap) |
| |
| global atexit_done |
| if not atexit_done: |
| atexit.register(reset_all) |
| atexit_done = True |
| |
| |
| def deinit(): |
| if orig_stdout is not None: |
| sys.stdout = orig_stdout |
| if orig_stderr is not None: |
| sys.stderr = orig_stderr |
| |
| |
| @contextlib.contextmanager |
| def colorama_text(*args, **kwargs): |
| init(*args, **kwargs) |
| try: |
| yield |
| finally: |
| deinit() |
| |
| |
| def reinit(): |
| if wrapped_stdout is not None: |
| sys.stdout = wrapped_stdout |
| if wrapped_stderr is not None: |
| sys.stderr = wrapped_stderr |
| |
| |
| def wrap_stream(stream, convert, strip, autoreset, wrap): |
| if wrap: |
| wrapper = AnsiToWin32(stream, |
| convert=convert, strip=strip, autoreset=autoreset) |
| if wrapper.should_wrap(): |
| stream = wrapper.stream |
| return stream |