Aluode/PerceptionLabPortable

1# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.2import atexit3import contextlib4import sys5 6from .ansitowin32 import AnsiToWin327 8 9def _wipe_internal_state_for_tests():10    global orig_stdout, orig_stderr11    orig_stdout = None12    orig_stderr = None13 14    global wrapped_stdout, wrapped_stderr15    wrapped_stdout = None16    wrapped_stderr = None17 18    global atexit_done19    atexit_done = False20 21    global fixed_windows_console22    fixed_windows_console = False23 24    try:25        # no-op if it wasn't registered26        atexit.unregister(reset_all)27    except AttributeError:28        # python 2: no atexit.unregister. Oh well, we did our best.29        pass30 31 32def reset_all():33    if AnsiToWin32 is not None:    # Issue #74: objects might become None at exit34        AnsiToWin32(orig_stdout).reset_all()35 36 37def init(autoreset=False, convert=None, strip=None, wrap=True):38 39    if not wrap and any([autoreset, convert, strip]):40        raise ValueError('wrap=False conflicts with any other arg=True')41 42    global wrapped_stdout, wrapped_stderr43    global orig_stdout, orig_stderr44 45    orig_stdout = sys.stdout46    orig_stderr = sys.stderr47 48    if sys.stdout is None:49        wrapped_stdout = None50    else:51        sys.stdout = wrapped_stdout = \52            wrap_stream(orig_stdout, convert, strip, autoreset, wrap)53    if sys.stderr is None:54        wrapped_stderr = None55    else:56        sys.stderr = wrapped_stderr = \57            wrap_stream(orig_stderr, convert, strip, autoreset, wrap)58 59    global atexit_done60    if not atexit_done:61        atexit.register(reset_all)62        atexit_done = True63 64 65def deinit():66    if orig_stdout is not None:67        sys.stdout = orig_stdout68    if orig_stderr is not None:69        sys.stderr = orig_stderr70 71 72def just_fix_windows_console():73    global fixed_windows_console74 75    if sys.platform != "win32":76        return77    if fixed_windows_console:78        return79    if wrapped_stdout is not None or wrapped_stderr is not None:80        # Someone already ran init() and it did stuff, so we won't second-guess them81        return82 83    # On newer versions of Windows, AnsiToWin32.__init__ will implicitly enable the84    # native ANSI support in the console as a side-effect. We only need to actually85    # replace sys.stdout/stderr if we're in the old-style conversion mode.86    new_stdout = AnsiToWin32(sys.stdout, convert=None, strip=None, autoreset=False)87    if new_stdout.convert:88        sys.stdout = new_stdout89    new_stderr = AnsiToWin32(sys.stderr, convert=None, strip=None, autoreset=False)90    if new_stderr.convert:91        sys.stderr = new_stderr92 93    fixed_windows_console = True94 95@contextlib.contextmanager96def colorama_text(*args, **kwargs):97    init(*args, **kwargs)98    try:99        yield100    finally:101        deinit()102 103 104def reinit():105    if wrapped_stdout is not None:106        sys.stdout = wrapped_stdout107    if wrapped_stderr is not None:108        sys.stderr = wrapped_stderr109 110 111def wrap_stream(stream, convert, strip, autoreset, wrap):112    if wrap:113        wrapper = AnsiToWin32(stream,114            convert=convert, strip=strip, autoreset=autoreset)115        if wrapper.should_wrap():116            stream = wrapper.stream117    return stream118 119 120# Use this for initial setup as well, to reduce code duplication121_wipe_internal_state_for_tests()122