|
|
# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. import atexit import contextlib import sys
from .ansitowin32 import AnsiToWin32
def _wipe_internal_state_for_tests(): global orig_stdout, orig_stderr orig_stdout = None orig_stderr = None
global wrapped_stdout, wrapped_stderr wrapped_stdout = None wrapped_stderr = None
global atexit_done atexit_done = False
global fixed_windows_console fixed_windows_console = False
try: # no-op if it wasn't registered atexit.unregister(reset_all) except AttributeError: # python 2: no atexit.unregister. Oh well, we did our best. pass
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
def just_fix_windows_console(): global fixed_windows_console
if sys.platform != "win32": return if fixed_windows_console: return if wrapped_stdout is not None or wrapped_stderr is not None: # Someone already ran init() and it did stuff, so we won't second-guess them return
# On newer versions of Windows, AnsiToWin32.__init__ will implicitly enable the # native ANSI support in the console as a side-effect. We only need to actually # replace sys.stdout/stderr if we're in the old-style conversion mode. new_stdout = AnsiToWin32(sys.stdout, convert=None, strip=None, autoreset=False) if new_stdout.convert: sys.stdout = new_stdout new_stderr = AnsiToWin32(sys.stderr, convert=None, strip=None, autoreset=False) if new_stderr.convert: sys.stderr = new_stderr
fixed_windows_console = True
@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
# Use this for initial setup as well, to reduce code duplication _wipe_internal_state_for_tests()
|