src/colorama/initialise.py

   1 # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
   2 import atexit
   3 import contextlib
   4 import sys
   5
   6 from .ansitowin32 import AnsiToWin32
   7
   8
   9 orig_stdout = None
  10 orig_stderr = None
  11
  12 wrapped_stdout = None
  13 wrapped_stderr = None
  14
  15 atexit_done = False
  16
  17
  18 def reset_all():
  19     if AnsiToWin32 is not None:  # Issue #74: objects might become None at exit
  20         AnsiToWin32(orig_stdout).reset_all()
  21
  22
  23 def init(autoreset=False, convert=None, strip=None, wrap=True):
  24
  25     if not wrap and any([autoreset, convert, strip]):
  26         raise ValueError("wrap=False conflicts with any other arg=True")
  27
  28     global wrapped_stdout, wrapped_stderr
  29     global orig_stdout, orig_stderr
  30
  31     orig_stdout = sys.stdout
  32     orig_stderr = sys.stderr
  33
  34     if sys.stdout is None:
  35         wrapped_stdout = None
  36     else:
  37         sys.stdout = wrapped_stdout = wrap_stream(
  38             orig_stdout, convert, strip, autoreset, wrap
  39         )
  40     if sys.stderr is None:
  41         wrapped_stderr = None
  42     else:
  43         sys.stderr = wrapped_stderr = wrap_stream(
  44             orig_stderr, convert, strip, autoreset, wrap
  45         )
  46
  47     global atexit_done
  48     if not atexit_done:
  49         atexit.register(reset_all)
  50         atexit_done = True
  51
  52
  53 def deinit():
  54     if orig_stdout is not None:
  55         sys.stdout = orig_stdout
  56     if orig_stderr is not None:
  57         sys.stderr = orig_stderr
  58
  59
  60 @contextlib.contextmanager
  61 def colorama_text(*args, **kwargs):
  62     init(*args, **kwargs)
  63     try:
  64         yield
  65     finally:
  66         deinit()
  67
  68
  69 def reinit():
  70     if wrapped_stdout is not None:
  71         sys.stdout = wrapped_stdout
  72     if wrapped_stderr is not None:
  73         sys.stderr = wrapped_stderr
  74
  75
  76 def wrap_stream(stream, convert, strip, autoreset, wrap):
  77     if wrap:
  78         wrapper = AnsiToWin32(stream, convert=convert, strip=strip, autoreset=autoreset)
  79         if wrapper.should_wrap():
  80             stream = wrapper.stream
  81     return stream