You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

initialise.py 1.9KB

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