diff --git a/colorama/initialise.py b/colorama/initialise.py index 6e01026..5074fd7 100644 --- a/colorama/initialise.py +++ b/colorama/initialise.py @@ -38,8 +38,10 @@ 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 + if orig_stdout is None: + orig_stdout = sys.stdout + if orig_stderr is None: + orig_stderr = sys.stderr if sys.stdout is None: wrapped_stdout = None @@ -54,7 +56,7 @@ def init(autoreset=False, convert=None, strip=None, wrap=True): global atexit_done if not atexit_done: - atexit.register(reset_all) + atexit.register(deinit) atexit_done = True diff --git a/tests/test_initialise.py b/tests/test_initialise.py new file mode 100644 index 0000000..1a1e758 --- /dev/null +++ b/tests/test_initialise.py @@ -0,0 +1,31 @@ +import sys +import unittest +from unittest.mock import patch, MagicMock +from colorama import init, deinit + +class TestColoramaInitialise(unittest.TestCase): + + @patch('colorama.ansitowin32.AnsiToWin32') + def test_init_deinit(self, mock_ansitowin32): + original_stdout = sys.stdout + original_stderr = sys.stderr + + # First call to init + init() + self.assertIsNotNone(sys.stdout) + self.assertIsNotNone(sys.stderr) + + # Second call to init + init() + self.assertIsNotNone(sys.stdout) + self.assertIsNotNone(sys.stderr) + + # Call deinit + deinit() + + # Verify stdout and stderr are restored correctly + self.assertIs(sys.stdout, original_stdout) + self.assertIs(sys.stderr, original_stderr) + +if __name__ == '__main__': + unittest.main() \ No newline at end of file