1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134
|
import contextlib
import io
import sys
import unittest
import unittest.mock
import _colorize
from test.support.os_helper import EnvironmentVarGuard
@contextlib.contextmanager
def clear_env():
with EnvironmentVarGuard() as mock_env:
mock_env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS", "TERM")
yield mock_env
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 test_colorized_detection_checks_for_environment_variables(self):
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,
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({'FORCE_COLOR': ''}, fallback, fallback)
check({'NO_COLOR': '1'}, fallback, False)
check({'NO_COLOR': '0'}, fallback, False)
check({'NO_COLOR': ''}, fallback, fallback)
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
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
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__":
unittest.main()
|