Improve tests for _colorize.can_colorize()

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__":