Updating colorama (Issue #1784)

thirdparty/colorama/initialise.py

@@ -1,20 +1,23 @@
+# Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file.
 import atexit
+import contextlib
 import sys
 
 from .ansitowin32 import AnsiToWin32
 
 
-orig_stdout = sys.stdout
-orig_stderr = sys.stderr
+orig_stdout = None
+orig_stderr = None
 
-wrapped_stdout = sys.stdout
-wrapped_stderr = sys.stderr
+wrapped_stdout = None
+wrapped_stderr = None
 
 atexit_done = False
 
 
 def reset_all():
-    AnsiToWin32(orig_stdout).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):
@@ -23,10 +26,21 @@ def init(autoreset=False, convert=None, strip=None, wrap=True):
         raise ValueError('wrap=False conflicts with any other arg=True')
 
     global wrapped_stdout, wrapped_stderr
-    sys.stdout = wrapped_stdout = \
-        wrap_stream(orig_stdout, convert, strip, autoreset, wrap)
-    sys.stderr = wrapped_stderr = \
-        wrap_stream(orig_stderr, convert, strip, autoreset, wrap)
+    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:
@@ -35,13 +49,26 @@ def init(autoreset=False, convert=None, strip=None, wrap=True):
 
 
 def deinit():
-    sys.stdout = orig_stdout
-    sys.stderr = orig_stderr
+    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():
-    sys.stdout = wrapped_stdout
-    sys.stderr = wrapped_stdout
+    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):