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 = \
  38             wrap_stream(orig_stdout, convert, strip, autoreset, wrap)
  39     if sys.stderr is None:
  40         wrapped_stderr = None
  41     else:
  42         sys.stderr = wrapped_stderr = \
  43             wrap_stream(orig_stderr, convert, strip, autoreset, wrap)
  44
  45     global atexit_done
  46     if not atexit_done:
  47         atexit.register(reset_all)
  48         atexit_done = True
  49
  50
  51 def deinit():
  52     if orig_stdout is not None:
  53         sys.stdout = orig_stdout
  54     if orig_stderr is not None:
  55         sys.stderr = orig_stderr
  56
  57
  58 @contextlib.contextmanager
  59 def colorama_text(*args, **kwargs):
  60     init(*args, **kwargs)
  61     try:
  62         yield
  63     finally:
  64         deinit()
  65
  66
  67 def reinit():
  68     if wrapped_stdout is not None:
  69         sys.stdout = wrapped_stdout
  70     if wrapped_stderr is not None:
  71         sys.stderr = wrapped_stderr
  72
  73
  74 def wrap_stream(stream, convert, strip, autoreset, wrap):
  75     if wrap:
  76         wrapper = AnsiToWin32(stream,
  77             convert=convert, strip=strip, autoreset=autoreset)
  78         if wrapper.should_wrap():
  79             stream = wrapper.stream
  80     return stream
  81
  82