You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
122 lines
3.2 KiB
Python
122 lines
3.2 KiB
Python
![]()
2 months ago
|
# 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()
|