|
|
- # 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
-
-
|