Spaces:
Runtime error
Runtime error
| # 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 | |
| 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() | |