diff --git a/colorama/initialise.py b/colorama/initialise.py index 4b45ba1..6e01026 100644 --- a/colorama/initialise.py +++ b/colorama/initialise.py @@ -6,15 +6,24 @@ from .ansitowin32 import AnsiToWin32 -orig_stdout = None -orig_stderr = None +def _wipe_internal_state_for_tests(): + global orig_stdout, orig_stderr + orig_stdout = None + orig_stderr = None + + global wrapped_stdout, wrapped_stderr + wrapped_stdout = None + wrapped_stderr = None + + global atexit_done + atexit_done = False -wrapped_stdout = None -wrapped_stderr = None + global fixed_windows_console + fixed_windows_console = False -atexit_done = False + # no-op if it wasn't registered + atexit.unregister(reset_all) -fixed_windows_console = False def reset_all(): if AnsiToWin32 is not None: # Issue #74: objects might become None at exit @@ -102,3 +111,7 @@ def wrap_stream(stream, convert, strip, autoreset, wrap): if wrapper.should_wrap(): stream = wrapper.stream return stream + + +# Use this for initial setup as well, to reduce code duplication +_wipe_internal_state_for_tests() diff --git a/colorama/tests/initialise_test.py b/colorama/tests/initialise_test.py index d689b4d..563985a 100644 --- a/colorama/tests/initialise_test.py +++ b/colorama/tests/initialise_test.py @@ -8,7 +8,7 @@ from mock import patch from ..ansitowin32 import StreamWrapper -from ..initialise import init +from ..initialise import init, just_fix_windows_console, _wipe_internal_state_for_tests from .utils import osname, replace_by orig_stdout = sys.stdout @@ -23,6 +23,7 @@ def setUp(self): self.assertNotWrapped() def tearDown(self): + _wipe_internal_state_for_tests() sys.stdout = orig_stdout sys.stderr = orig_stderr