diff --git a/Lib/test/test__colorize.py b/Lib/test/test__colorize.py index 25519ba7e92e1f..ff15d83cb078fe 100644 --- a/Lib/test/test__colorize.py +++ b/Lib/test/test__colorize.py @@ -1,80 +1,132 @@ import contextlib +import io import sys import unittest import unittest.mock import _colorize -from test.support import force_not_colorized, make_clean_env +from test.support.os_helper import EnvironmentVarGuard -ORIGINAL_CAN_COLORIZE = _colorize.can_colorize +@contextlib.contextmanager +def clear_env(): + with EnvironmentVarGuard() as mock_env: + for var in "FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS": + mock_env.unset(var) + yield mock_env -def setUpModule(): - _colorize.can_colorize = lambda *args, **kwargs: False - -def tearDownModule(): - _colorize.can_colorize = ORIGINAL_CAN_COLORIZE +def supports_virtual_terminal(): + if sys.platform == "win32": + return unittest.mock.patch("nt._supports_virtual_terminal", return_value=True) + else: + return contextlib.nullcontext() class TestColorizeFunction(unittest.TestCase): - def setUp(self): - # Remove PYTHON* environment variables to isolate from local user - # settings and simulate running with `-E`. Such variables should be - # added to test methods later to patched os.environ. - patcher = unittest.mock.patch("os.environ", new=make_clean_env()) - self.addCleanup(patcher.stop) - patcher.start() - - @force_not_colorized def test_colorized_detection_checks_for_environment_variables(self): - flags = unittest.mock.MagicMock(ignore_environment=False) + def check(env, fallback, expected): + with (self.subTest(env=env, fallback=fallback), + clear_env() as mock_env): + mock_env.update(env) + isatty_mock.return_value = fallback + stdout_mock.isatty.return_value = fallback + self.assertEqual(_colorize.can_colorize(), expected) + with (unittest.mock.patch("os.isatty") as isatty_mock, unittest.mock.patch("sys.stdout") as stdout_mock, - unittest.mock.patch("sys.stderr") as stderr_mock, - unittest.mock.patch("sys.flags", flags), - unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE), - (unittest.mock.patch("nt._supports_virtual_terminal", return_value=False) - if sys.platform == "win32" else - contextlib.nullcontext()) as vt_mock): + supports_virtual_terminal()): + stdout_mock.fileno.return_value = 1 + + for fallback in False, True: + check({}, fallback, fallback) + check({'TERM': 'dumb'}, fallback, False) + check({'TERM': 'xterm'}, fallback, fallback) + check({'TERM': ''}, fallback, fallback) + check({'FORCE_COLOR': '1'}, fallback, True) + check({'FORCE_COLOR': '0'}, fallback, True) + check({'NO_COLOR': '1'}, fallback, False) + check({'NO_COLOR': '0'}, fallback, False) + + check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True) + check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False) + for ignore_environment in False, True: + # Simulate running with or without `-E`. + flags = unittest.mock.MagicMock(ignore_environment=ignore_environment) + with unittest.mock.patch("sys.flags", flags): + check({'PYTHON_COLORS': '1'}, True, True) + check({'PYTHON_COLORS': '1'}, False, not ignore_environment) + check({'PYTHON_COLORS': '0'}, True, ignore_environment) + check({'PYTHON_COLORS': '0'}, False, False) + for fallback in False, True: + check({'PYTHON_COLORS': 'x'}, fallback, fallback) + check({'PYTHON_COLORS': ''}, fallback, fallback) + + check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment) + check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment) + check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment) + + @unittest.skipUnless(sys.platform == "win32", "requires Windows") + def test_colorized_detection_checks_on_windows(self): + with (clear_env(), + unittest.mock.patch("os.isatty") as isatty_mock, + unittest.mock.patch("sys.stdout") as stdout_mock, + supports_virtual_terminal() as vt_mock): + stdout_mock.fileno.return_value = 1 isatty_mock.return_value = True + stdout_mock.isatty.return_value = True + + vt_mock.return_value = True + self.assertEqual(_colorize.can_colorize(), True) + vt_mock.return_value = False + self.assertEqual(_colorize.can_colorize(), False) + import nt + del nt._supports_virtual_terminal + self.assertEqual(_colorize.can_colorize(), False) + + def test_colorized_detection_checks_for_std_streams(self): + with (clear_env(), + unittest.mock.patch("os.isatty") as isatty_mock, + unittest.mock.patch("sys.stdout") as stdout_mock, + unittest.mock.patch("sys.stderr") as stderr_mock, + supports_virtual_terminal()): stdout_mock.fileno.return_value = 1 + stderr_mock.fileno.side_effect = ZeroDivisionError + stderr_mock.isatty.side_effect = ZeroDivisionError + + isatty_mock.return_value = True stdout_mock.isatty.return_value = True - stderr_mock.fileno.return_value = 2 - stderr_mock.isatty.return_value = True - with unittest.mock.patch("os.environ", {'TERM': 'dumb'}): - self.assertEqual(_colorize.can_colorize(), False) - with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}): - self.assertEqual(_colorize.can_colorize(), True) - with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}): - self.assertEqual(_colorize.can_colorize(), False) - with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}): - self.assertEqual(_colorize.can_colorize(), False) - with unittest.mock.patch("os.environ", - {'NO_COLOR': '1', "PYTHON_COLORS": '1'}): - self.assertEqual(_colorize.can_colorize(), True) - with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}): - self.assertEqual(_colorize.can_colorize(), True) - with unittest.mock.patch("os.environ", - {'FORCE_COLOR': '1', 'NO_COLOR': '1'}): - self.assertEqual(_colorize.can_colorize(), False) - with unittest.mock.patch("os.environ", - {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}): - self.assertEqual(_colorize.can_colorize(), False) - - with unittest.mock.patch("os.environ", {}): - if sys.platform == "win32": - self.assertEqual(_colorize.can_colorize(), False) - - vt_mock.return_value = True - self.assertEqual(_colorize.can_colorize(), True) - else: - self.assertEqual(_colorize.can_colorize(), True) + self.assertEqual(_colorize.can_colorize(), True) + + isatty_mock.return_value = False + stdout_mock.isatty.return_value = False + self.assertEqual(_colorize.can_colorize(), False) + + def test_colorized_detection_checks_for_file(self): + with clear_env(), supports_virtual_terminal(): + with unittest.mock.patch("os.isatty") as isatty_mock: + file = unittest.mock.MagicMock() + file.fileno.return_value = 1 + isatty_mock.return_value = True + self.assertEqual(_colorize.can_colorize(file=file), True) isatty_mock.return_value = False - stdout_mock.isatty.return_value = False - stderr_mock.isatty.return_value = False - self.assertEqual(_colorize.can_colorize(), False) + self.assertEqual(_colorize.can_colorize(file=file), False) + + # No file.fileno. + with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError): + file = unittest.mock.MagicMock(spec=['isatty']) + file.isatty.return_value = True + self.assertEqual(_colorize.can_colorize(file=file), False) + + # file.fileno() raises io.UnsupportedOperation. + with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError): + file = unittest.mock.MagicMock() + file.fileno.side_effect = io.UnsupportedOperation + file.isatty.return_value = True + self.assertEqual(_colorize.can_colorize(file=file), True) + file.isatty.return_value = False + self.assertEqual(_colorize.can_colorize(file=file), False) if __name__ == "__main__":