Prevent user from wrapping already wrapped streams

tartley · Nov 6, 2018 · 9bbe4e8 · 9bbe4e8
1 parent c11da6f
commit 9bbe4e8
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/colorama/initialise.py b/colorama/initialise.py
@@ -28,6 +28,9 @@ def init(autoreset=False, convert=None, strip=None, wrap=True):
     global wrapped_stdout, wrapped_stderr
     global orig_stdout, orig_stderr
 
+    if (orig_stdout is not None) or (orig_stderr is not None):
+        deinit()
+
     orig_stdout = sys.stdout
     orig_stderr = sys.stderr
 
@@ -49,10 +52,14 @@ def init(autoreset=False, convert=None, strip=None, wrap=True):
 
 
 def deinit():
+    global orig_stdout, orig_stderr
+
     if orig_stdout is not None:
         sys.stdout = orig_stdout
+        orig_stdout = None
     if orig_stderr is not None:
         sys.stderr = orig_stderr
+        orig_stderr = None
 
 
 @contextlib.contextmanager

diff --git a/colorama/tests/initialise_test.py b/colorama/tests/initialise_test.py
@@ -6,7 +6,7 @@
 from mock import patch
 
 from ..ansitowin32 import StreamWrapper
-from ..initialise import init
+from ..initialise import init, deinit
 from .utils import osname, redirected_output, replace_by
 
 orig_stdout = sys.stdout
@@ -22,6 +22,7 @@ def setUp(self):
     def tearDown(self):
         sys.stdout = orig_stdout
         sys.stderr = orig_stderr
+        deinit()
 
     def assertWrapped(self):
         self.assertIsNot(sys.stdout, orig_stdout, 'stdout should be wrapped')