From 0d55855b32ee27aa931399213f4d3655f65be9f7 Mon Sep 17 00:00:00 2001 From: Jonathan Hartley Date: Mon, 11 Nov 2019 10:07:13 -0600 Subject: [PATCH 1/3] Guard against multiple calls to init(). BEWARE, I haven't tested this on Windows (because I don't have a Windows box.) It should definitely get tested there before we merge. This should no longer lose the original stdout/stderr, so that we can restore them with a single call to deinit(), while still preserving the behavior that multiple calls to init() can override the passed values for autoreset, convert, strip, wrap. This fix is not intended to reset output to default foreground/background colors after the call to 'deinit()' - that was never intended behavior. --- colorama/initialise.py | 22 ++++++++++++++++------ colorama/tests/initialise_test.py | 12 +++++++----- 2 files changed, 23 insertions(+), 11 deletions(-) diff --git a/colorama/initialise.py b/colorama/initialise.py index 430d0668..7abeb1cf 100644 --- a/colorama/initialise.py +++ b/colorama/initialise.py @@ -5,9 +5,10 @@ from .ansitowin32 import AnsiToWin32 +UNSET = object() -orig_stdout = None -orig_stderr = None +orig_stdout = UNSET +orig_stderr = UNSET wrapped_stdout = None wrapped_stderr = None @@ -28,8 +29,13 @@ def init(autoreset=False, convert=None, strip=None, wrap=True): global wrapped_stdout, wrapped_stderr global orig_stdout, orig_stderr - orig_stdout = sys.stdout - orig_stderr = sys.stderr + # Prevent multiple calls from losing the original stdout + if ( + orig_stdout is UNSET and + orig_stderr is UNSET + ): + orig_stdout = sys.stdout + orig_stderr = sys.stderr if sys.stdout is None: wrapped_stdout = None @@ -49,10 +55,13 @@ def init(autoreset=False, convert=None, strip=None, wrap=True): def deinit(): - if orig_stdout is not None: + global orig_stdout, orig_stderr + if orig_stdout is not UNSET: sys.stdout = orig_stdout - if orig_stderr is not None: + orig_stdout == UNSET + if orig_stderr is not UNSET: sys.stderr = orig_stderr + orig_stderr == UNSET @contextlib.contextmanager @@ -78,3 +87,4 @@ def wrap_stream(stream, convert, strip, autoreset, wrap): if wrapper.should_wrap(): stream = wrapper.stream return stream + diff --git a/colorama/tests/initialise_test.py b/colorama/tests/initialise_test.py index 2f7384de..24ad7448 100644 --- a/colorama/tests/initialise_test.py +++ b/colorama/tests/initialise_test.py @@ -6,13 +6,12 @@ from mock import patch from ..ansitowin32 import StreamWrapper -from ..initialise import init +from ..initialise import deinit, init from .utils import osname, redirected_output, replace_by orig_stdout = sys.stdout orig_stderr = sys.stderr - class InitTest(TestCase): def setUp(self): @@ -75,13 +74,16 @@ def testInitWrapOffDoesntWrapOnWindows(self): def testInitWrapOffIncompatibleWithAutoresetOn(self): self.assertRaises(ValueError, lambda: init(autoreset=True, wrap=False)) - @patch('colorama.ansitowin32.winterm', None) @patch('colorama.ansitowin32.winapi_test', lambda *_: True) - def testInitOnlyWrapsOnce(self): - with osname("nt"): + def testInitTwiceCanUndoneWithDeinitOnce(self): + with osname('nt'): + self.assertNotWrapped() init() + self.assertWrapped() init() self.assertWrapped() + deinit() + self.assertNotWrapped() @patch('colorama.win32.SetConsoleTextAttribute') @patch('colorama.initialise.AnsiToWin32') From 66decc2b8ea34bac5bbb2e7871c34c9680f8f87b Mon Sep 17 00:00:00 2001 From: Jonathan Hartley Date: Mon, 11 Nov 2019 11:30:38 -0600 Subject: [PATCH 2/3] test context managers should clean up in ...finally --- colorama/tests/utils.py | 37 ++++++++++++++++++++++++------------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/colorama/tests/utils.py b/colorama/tests/utils.py index de2abf55..a8af4419 100644 --- a/colorama/tests/utils.py +++ b/colorama/tests/utils.py @@ -18,16 +18,20 @@ def isatty(self): def osname(name): orig = os.name os.name = name - yield - os.name = orig + try: + yield + finally: + os.name = orig @contextmanager def redirected_output(): orig = sys.stdout sys.stdout = Mock() sys.stdout.isatty = lambda: False - yield - sys.stdout = orig + try: + yield + finally: + sys.stdout = orig @contextmanager def replace_by(stream): @@ -35,9 +39,11 @@ def replace_by(stream): orig_stderr = sys.stderr sys.stdout = stream sys.stderr = stream - yield - sys.stdout = orig_stdout - sys.stderr = orig_stderr + try: + yield + finally: + sys.stdout = orig_stdout + sys.stderr = orig_stderr @contextmanager def replace_original_by(stream): @@ -45,14 +51,19 @@ def replace_original_by(stream): orig_stderr = sys.__stderr__ sys.__stdout__ = stream sys.__stderr__ = stream - yield - sys.__stdout__ = orig_stdout - sys.__stderr__ = orig_stderr + try: + yield + finally: + sys.__stdout__ = orig_stdout + sys.__stderr__ = orig_stderr @contextmanager def pycharm(): os.environ["PYCHARM_HOSTED"] = "1" non_tty = StreamNonTTY() - with replace_by(non_tty), replace_original_by(non_tty): - yield - del os.environ["PYCHARM_HOSTED"] + try: + with replace_by(non_tty), replace_original_by(non_tty): + yield + finally: + del os.environ["PYCHARM_HOSTED"] + From 16d462b115a6a2d3a51f3b0aad2c5cec1d8c300a Mon Sep 17 00:00:00 2001 From: Jonathan Hartley Date: Mon, 11 Nov 2019 11:38:19 -0600 Subject: [PATCH 3/3] Fix broken assignment in deinit. Add an extra test to prove that it's needed. --- colorama/initialise.py | 13 +++++-------- colorama/tests/initialise_test.py | 21 ++++++++++++++++++++- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/colorama/initialise.py b/colorama/initialise.py index 7abeb1cf..eb114b96 100644 --- a/colorama/initialise.py +++ b/colorama/initialise.py @@ -22,19 +22,16 @@ def reset_all(): def init(autoreset=False, convert=None, strip=None, wrap=True): - if not wrap and any([autoreset, convert, strip]): raise ValueError('wrap=False conflicts with any other arg=True') global wrapped_stdout, wrapped_stderr global orig_stdout, orig_stderr - # Prevent multiple calls from losing the original stdout - if ( - orig_stdout is UNSET and - orig_stderr is UNSET - ): + # Prevent multiple calls from losing the original stdout/err + if orig_stdout is UNSET: orig_stdout = sys.stdout + if orig_stderr is UNSET: orig_stderr = sys.stderr if sys.stdout is None: @@ -58,10 +55,10 @@ def deinit(): global orig_stdout, orig_stderr if orig_stdout is not UNSET: sys.stdout = orig_stdout - orig_stdout == UNSET + orig_stdout = UNSET if orig_stderr is not UNSET: sys.stderr = orig_stderr - orig_stderr == UNSET + orig_stderr = UNSET @contextlib.contextmanager diff --git a/colorama/tests/initialise_test.py b/colorama/tests/initialise_test.py index 24ad7448..0aac36ec 100644 --- a/colorama/tests/initialise_test.py +++ b/colorama/tests/initialise_test.py @@ -6,12 +6,14 @@ from mock import patch from ..ansitowin32 import StreamWrapper +from .. import initialise from ..initialise import deinit, init from .utils import osname, redirected_output, replace_by orig_stdout = sys.stdout orig_stderr = sys.stderr + class InitTest(TestCase): def setUp(self): @@ -21,6 +23,8 @@ def setUp(self): def tearDown(self): sys.stdout = orig_stdout sys.stderr = orig_stderr + initialise.orig_stdout = initialise.UNSET + initialise.orig_stderr = initialise.UNSET def assertWrapped(self): self.assertIsNot(sys.stdout, orig_stdout, 'stdout should be wrapped') @@ -75,7 +79,7 @@ def testInitWrapOffIncompatibleWithAutoresetOn(self): self.assertRaises(ValueError, lambda: init(autoreset=True, wrap=False)) @patch('colorama.ansitowin32.winapi_test', lambda *_: True) - def testInitTwiceCanUndoneWithDeinitOnce(self): + def testInitTwiceCanBeUndoneWithDeinitOnce(self): with osname('nt'): self.assertNotWrapped() init() @@ -85,6 +89,21 @@ def testInitTwiceCanUndoneWithDeinitOnce(self): deinit() self.assertNotWrapped() + @patch('colorama.ansitowin32.winapi_test', lambda *_: True) + def testInitDeinitInitWorks(self): + with osname('nt'): + self.assertNotWrapped() + init() + self.assertWrapped() + deinit() + self.assertNotWrapped() + init() + self.assertWrapped() + deinit() + self.assertNotWrapped() + init() + self.assertWrapped() + @patch('colorama.win32.SetConsoleTextAttribute') @patch('colorama.initialise.AnsiToWin32') def testAutoResetPassedOn(self, mockATW32, _):