|  | 
| 1 | 1 | import contextlib | 
|  | 2 | +import io | 
| 2 | 3 | import sys | 
| 3 | 4 | import unittest | 
| 4 | 5 | import unittest.mock | 
| 5 | 6 | import _colorize | 
| 6 |  | -from test.support import force_not_colorized, make_clean_env | 
|  | 7 | +from test.support.os_helper import EnvironmentVarGuard | 
| 7 | 8 | 
 | 
| 8 |  | -ORIGINAL_CAN_COLORIZE = _colorize.can_colorize | 
| 9 | 9 | 
 | 
|  | 10 | +@contextlib.contextmanager | 
|  | 11 | +def clear_env(): | 
|  | 12 | +    with EnvironmentVarGuard() as mock_env: | 
|  | 13 | +        for var in "FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS": | 
|  | 14 | +            mock_env.unset(var) | 
|  | 15 | +        yield mock_env | 
| 10 | 16 | 
 | 
| 11 |  | -def setUpModule(): | 
| 12 |  | -    _colorize.can_colorize = lambda *args, **kwargs: False | 
| 13 | 17 | 
 | 
| 14 |  | - | 
| 15 |  | -def tearDownModule(): | 
| 16 |  | -    _colorize.can_colorize = ORIGINAL_CAN_COLORIZE | 
|  | 18 | +def supports_virtual_terminal(): | 
|  | 19 | +    if sys.platform == "win32": | 
|  | 20 | +        return unittest.mock.patch("nt._supports_virtual_terminal", return_value=True) | 
|  | 21 | +    else: | 
|  | 22 | +        return contextlib.nullcontext() | 
| 17 | 23 | 
 | 
| 18 | 24 | 
 | 
| 19 | 25 | class TestColorizeFunction(unittest.TestCase): | 
| 20 |  | -    def setUp(self): | 
| 21 |  | -        # Remove PYTHON* environment variables to isolate from local user | 
| 22 |  | -        # settings and simulate running with `-E`. Such variables should be | 
| 23 |  | -        # added to test methods later to patched os.environ. | 
| 24 |  | -        patcher = unittest.mock.patch("os.environ", new=make_clean_env()) | 
| 25 |  | -        self.addCleanup(patcher.stop) | 
| 26 |  | -        patcher.start() | 
| 27 |  | - | 
| 28 |  | -    @force_not_colorized | 
| 29 | 26 |     def test_colorized_detection_checks_for_environment_variables(self): | 
| 30 |  | -        flags = unittest.mock.MagicMock(ignore_environment=False) | 
|  | 27 | +        def check(env, fallback, expected): | 
|  | 28 | +            with (self.subTest(env=env, fallback=fallback), | 
|  | 29 | +                  clear_env() as mock_env): | 
|  | 30 | +                mock_env.update(env) | 
|  | 31 | +                isatty_mock.return_value = fallback | 
|  | 32 | +                stdout_mock.isatty.return_value = fallback | 
|  | 33 | +                self.assertEqual(_colorize.can_colorize(), expected) | 
|  | 34 | + | 
| 31 | 35 |         with (unittest.mock.patch("os.isatty") as isatty_mock, | 
| 32 | 36 |               unittest.mock.patch("sys.stdout") as stdout_mock, | 
| 33 |  | -              unittest.mock.patch("sys.stderr") as stderr_mock, | 
| 34 |  | -              unittest.mock.patch("sys.flags", flags), | 
| 35 |  | -              unittest.mock.patch("_colorize.can_colorize", ORIGINAL_CAN_COLORIZE), | 
| 36 |  | -              (unittest.mock.patch("nt._supports_virtual_terminal", return_value=False) | 
| 37 |  | -               if sys.platform == "win32" else | 
| 38 |  | -               contextlib.nullcontext()) as vt_mock): | 
|  | 37 | +              supports_virtual_terminal()): | 
|  | 38 | +            stdout_mock.fileno.return_value = 1 | 
|  | 39 | + | 
|  | 40 | +            for fallback in False, True: | 
|  | 41 | +                check({}, fallback, fallback) | 
|  | 42 | +                check({'TERM': 'dumb'}, fallback, False) | 
|  | 43 | +                check({'TERM': 'xterm'}, fallback, fallback) | 
|  | 44 | +                check({'TERM': ''}, fallback, fallback) | 
|  | 45 | +                check({'FORCE_COLOR': '1'}, fallback, True) | 
|  | 46 | +                check({'FORCE_COLOR': '0'}, fallback, True) | 
|  | 47 | +                check({'NO_COLOR': '1'}, fallback, False) | 
|  | 48 | +                check({'NO_COLOR': '0'}, fallback, False) | 
|  | 49 | + | 
|  | 50 | +            check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True) | 
|  | 51 | +            check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False) | 
| 39 | 52 | 
 | 
|  | 53 | +            for ignore_environment in False, True: | 
|  | 54 | +                # Simulate running with or without `-E`. | 
|  | 55 | +                flags = unittest.mock.MagicMock(ignore_environment=ignore_environment) | 
|  | 56 | +                with unittest.mock.patch("sys.flags", flags): | 
|  | 57 | +                    check({'PYTHON_COLORS': '1'}, True, True) | 
|  | 58 | +                    check({'PYTHON_COLORS': '1'}, False, not ignore_environment) | 
|  | 59 | +                    check({'PYTHON_COLORS': '0'}, True, ignore_environment) | 
|  | 60 | +                    check({'PYTHON_COLORS': '0'}, False, False) | 
|  | 61 | +                    for fallback in False, True: | 
|  | 62 | +                        check({'PYTHON_COLORS': 'x'}, fallback, fallback) | 
|  | 63 | +                        check({'PYTHON_COLORS': ''}, fallback, fallback) | 
|  | 64 | + | 
|  | 65 | +                    check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment) | 
|  | 66 | +                    check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment) | 
|  | 67 | +                    check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment) | 
|  | 68 | + | 
|  | 69 | +    @unittest.skipUnless(sys.platform == "win32", "requires Windows") | 
|  | 70 | +    def test_colorized_detection_checks_on_windows(self): | 
|  | 71 | +        with (clear_env(), | 
|  | 72 | +              unittest.mock.patch("os.isatty") as isatty_mock, | 
|  | 73 | +              unittest.mock.patch("sys.stdout") as stdout_mock, | 
|  | 74 | +              supports_virtual_terminal() as vt_mock): | 
|  | 75 | +            stdout_mock.fileno.return_value = 1 | 
| 40 | 76 |             isatty_mock.return_value = True | 
|  | 77 | +            stdout_mock.isatty.return_value = True | 
|  | 78 | + | 
|  | 79 | +            vt_mock.return_value = True | 
|  | 80 | +            self.assertEqual(_colorize.can_colorize(), True) | 
|  | 81 | +            vt_mock.return_value = False | 
|  | 82 | +            self.assertEqual(_colorize.can_colorize(), False) | 
|  | 83 | +            import nt | 
|  | 84 | +            del nt._supports_virtual_terminal | 
|  | 85 | +            self.assertEqual(_colorize.can_colorize(), False) | 
|  | 86 | + | 
|  | 87 | +    def test_colorized_detection_checks_for_std_streams(self): | 
|  | 88 | +        with (clear_env(), | 
|  | 89 | +              unittest.mock.patch("os.isatty") as isatty_mock, | 
|  | 90 | +              unittest.mock.patch("sys.stdout") as stdout_mock, | 
|  | 91 | +              unittest.mock.patch("sys.stderr") as stderr_mock, | 
|  | 92 | +              supports_virtual_terminal()): | 
| 41 | 93 |             stdout_mock.fileno.return_value = 1 | 
|  | 94 | +            stderr_mock.fileno.side_effect = ZeroDivisionError | 
|  | 95 | +            stderr_mock.isatty.side_effect = ZeroDivisionError | 
|  | 96 | + | 
|  | 97 | +            isatty_mock.return_value = True | 
| 42 | 98 |             stdout_mock.isatty.return_value = True | 
| 43 |  | -            stderr_mock.fileno.return_value = 2 | 
| 44 |  | -            stderr_mock.isatty.return_value = True | 
| 45 |  | -            with unittest.mock.patch("os.environ", {'TERM': 'dumb'}): | 
| 46 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
| 47 |  | -            with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '1'}): | 
| 48 |  | -                self.assertEqual(_colorize.can_colorize(), True) | 
| 49 |  | -            with unittest.mock.patch("os.environ", {'PYTHON_COLORS': '0'}): | 
| 50 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
| 51 |  | -            with unittest.mock.patch("os.environ", {'NO_COLOR': '1'}): | 
| 52 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
| 53 |  | -            with unittest.mock.patch("os.environ", | 
| 54 |  | -                                     {'NO_COLOR': '1', "PYTHON_COLORS": '1'}): | 
| 55 |  | -                self.assertEqual(_colorize.can_colorize(), True) | 
| 56 |  | -            with unittest.mock.patch("os.environ", {'FORCE_COLOR': '1'}): | 
| 57 |  | -                self.assertEqual(_colorize.can_colorize(), True) | 
| 58 |  | -            with unittest.mock.patch("os.environ", | 
| 59 |  | -                                     {'FORCE_COLOR': '1', 'NO_COLOR': '1'}): | 
| 60 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
| 61 |  | -            with unittest.mock.patch("os.environ", | 
| 62 |  | -                                     {'FORCE_COLOR': '1', "PYTHON_COLORS": '0'}): | 
| 63 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
| 64 |  | - | 
| 65 |  | -            with unittest.mock.patch("os.environ", {}): | 
| 66 |  | -                if sys.platform == "win32": | 
| 67 |  | -                    self.assertEqual(_colorize.can_colorize(), False) | 
| 68 |  | - | 
| 69 |  | -                    vt_mock.return_value = True | 
| 70 |  | -                    self.assertEqual(_colorize.can_colorize(), True) | 
| 71 |  | -                else: | 
| 72 |  | -                    self.assertEqual(_colorize.can_colorize(), True) | 
|  | 99 | +            self.assertEqual(_colorize.can_colorize(), True) | 
|  | 100 | + | 
|  | 101 | +            isatty_mock.return_value = False | 
|  | 102 | +            stdout_mock.isatty.return_value = False | 
|  | 103 | +            self.assertEqual(_colorize.can_colorize(), False) | 
|  | 104 | + | 
|  | 105 | +    def test_colorized_detection_checks_for_file(self): | 
|  | 106 | +        with clear_env(), supports_virtual_terminal(): | 
| 73 | 107 | 
 | 
|  | 108 | +            with unittest.mock.patch("os.isatty") as isatty_mock: | 
|  | 109 | +                file = unittest.mock.MagicMock() | 
|  | 110 | +                file.fileno.return_value = 1 | 
|  | 111 | +                isatty_mock.return_value = True | 
|  | 112 | +                self.assertEqual(_colorize.can_colorize(file=file), True) | 
| 74 | 113 |                 isatty_mock.return_value = False | 
| 75 |  | -                stdout_mock.isatty.return_value = False | 
| 76 |  | -                stderr_mock.isatty.return_value = False | 
| 77 |  | -                self.assertEqual(_colorize.can_colorize(), False) | 
|  | 114 | +                self.assertEqual(_colorize.can_colorize(file=file), False) | 
|  | 115 | + | 
|  | 116 | +            # No file.fileno. | 
|  | 117 | +            with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError): | 
|  | 118 | +                file = unittest.mock.MagicMock(spec=['isatty']) | 
|  | 119 | +                file.isatty.return_value = True | 
|  | 120 | +                self.assertEqual(_colorize.can_colorize(file=file), False) | 
|  | 121 | + | 
|  | 122 | +            # file.fileno() raises io.UnsupportedOperation. | 
|  | 123 | +            with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError): | 
|  | 124 | +                file = unittest.mock.MagicMock() | 
|  | 125 | +                file.fileno.side_effect = io.UnsupportedOperation | 
|  | 126 | +                file.isatty.return_value = True | 
|  | 127 | +                self.assertEqual(_colorize.can_colorize(file=file), True) | 
|  | 128 | +                file.isatty.return_value = False | 
|  | 129 | +                self.assertEqual(_colorize.can_colorize(file=file), False) | 
| 78 | 130 | 
 | 
| 79 | 131 | 
 | 
| 80 | 132 | if __name__ == "__main__": | 
|  | 
0 commit comments