forked from Rust-related/RustPython
add _colorize.py at 3.13.2
Signed-off-by: Ashwin Naren <arihant2math@gmail.com>
This commit is contained in:
Ashwin Naren
67
Lib/_colorize.py
vendored
Normal file
67
Lib/_colorize.py
vendored
Normal file
@@ -0,0 +1,67 @@
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
|
||||
COLORIZE = True
|
||||
|
||||
|
||||
class ANSIColors:
|
||||
BOLD_GREEN = "\x1b[1;32m"
|
||||
BOLD_MAGENTA = "\x1b[1;35m"
|
||||
BOLD_RED = "\x1b[1;31m"
|
||||
GREEN = "\x1b[32m"
|
||||
GREY = "\x1b[90m"
|
||||
MAGENTA = "\x1b[35m"
|
||||
RED = "\x1b[31m"
|
||||
RESET = "\x1b[0m"
|
||||
YELLOW = "\x1b[33m"
|
||||
|
||||
|
||||
NoColors = ANSIColors()
|
||||
|
||||
for attr in dir(NoColors):
|
||||
if not attr.startswith("__"):
|
||||
setattr(NoColors, attr, "")
|
||||
|
||||
|
||||
def get_colors(colorize: bool = False, *, file=None) -> ANSIColors:
|
||||
if colorize or can_colorize(file=file):
|
||||
return ANSIColors()
|
||||
else:
|
||||
return NoColors
|
||||
|
||||
|
||||
def can_colorize(*, file=None) -> bool:
|
||||
if file is None:
|
||||
file = sys.stdout
|
||||
|
||||
if not sys.flags.ignore_environment:
|
||||
if os.environ.get("PYTHON_COLORS") == "0":
|
||||
return False
|
||||
if os.environ.get("PYTHON_COLORS") == "1":
|
||||
return True
|
||||
if os.environ.get("NO_COLOR"):
|
||||
return False
|
||||
if not COLORIZE:
|
||||
return False
|
||||
if os.environ.get("FORCE_COLOR"):
|
||||
return True
|
||||
if os.environ.get("TERM") == "dumb":
|
||||
return False
|
||||
|
||||
if not hasattr(file, "fileno"):
|
||||
return False
|
||||
|
||||
if sys.platform == "win32":
|
||||
try:
|
||||
import nt
|
||||
|
||||
if not nt._supports_virtual_terminal():
|
||||
return False
|
||||
except (ImportError, AttributeError):
|
||||
return False
|
||||
|
||||
try:
|
||||
return os.isatty(file.fileno())
|
||||
except io.UnsupportedOperation:
|
||||
return file.isatty()
|
||||
135
Lib/test/test__colorize.py
vendored
Normal file
135
Lib/test/test__colorize.py
vendored
Normal file
@@ -0,0 +1,135 @@
|
||||
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:
|
||||
for var in "FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS":
|
||||
mock_env.unset(var)
|
||||
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()
|
||||
Reference in New Issue
Block a user