Skip to content

Commit 19050af

Browse files
youknowoneyouknowone
authored
Merge pull request RustPython#5520 from arihant2math/colorize
Add _colorize at 3.13.2
2 parents a46ce8e + e96557b commit 19050af

File tree

3 files changed

+208
-0
lines changed

3 files changed

+208
-0
lines changed

‎Lib/_colorize.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
import io
2+
import os
3+
import sys
4+
5+
COLORIZE = True
6+
7+
8+
class ANSIColors:
9+
BOLD_GREEN = "\x1b[1;32m"
10+
BOLD_MAGENTA = "\x1b[1;35m"
11+
BOLD_RED = "\x1b[1;31m"
12+
GREEN = "\x1b[32m"
13+
GREY = "\x1b[90m"
14+
MAGENTA = "\x1b[35m"
15+
RED = "\x1b[31m"
16+
RESET = "\x1b[0m"
17+
YELLOW = "\x1b[33m"
18+
19+
20+
NoColors = ANSIColors()
21+
22+
for attr in dir(NoColors):
23+
if not attr.startswith("__"):
24+
setattr(NoColors, attr, "")
25+
26+
27+
def get_colors(colorize: bool = False, *, file=None) -> ANSIColors:
28+
if colorize or can_colorize(file=file):
29+
return ANSIColors()
30+
else:
31+
return NoColors
32+
33+
34+
def can_colorize(*, file=None) -> bool:
35+
if file is None:
36+
file = sys.stdout
37+
38+
if not sys.flags.ignore_environment:
39+
if os.environ.get("PYTHON_COLORS") == "0":
40+
return False
41+
if os.environ.get("PYTHON_COLORS") == "1":
42+
return True
43+
if os.environ.get("NO_COLOR"):
44+
return False
45+
if not COLORIZE:
46+
return False
47+
if os.environ.get("FORCE_COLOR"):
48+
return True
49+
if os.environ.get("TERM") == "dumb":
50+
return False
51+
52+
if not hasattr(file, "fileno"):
53+
return False
54+
55+
if sys.platform == "win32":
56+
try:
57+
import nt
58+
59+
if not nt._supports_virtual_terminal():
60+
return False
61+
except (ImportError, AttributeError):
62+
return False
63+
64+
try:
65+
return os.isatty(file.fileno())
66+
except io.UnsupportedOperation:
67+
return file.isatty()

‎Lib/test/test__colorize.py

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
import contextlib
2+
import io
3+
import sys
4+
import unittest
5+
import unittest.mock
6+
import _colorize
7+
from test.support.os_helper import EnvironmentVarGuard
8+
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
16+
17+
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()
23+
24+
25+
class TestColorizeFunction(unittest.TestCase):
26+
def test_colorized_detection_checks_for_environment_variables(self):
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+
35+
with (unittest.mock.patch("os.isatty") as isatty_mock,
36+
unittest.mock.patch("sys.stdout") as stdout_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({'FORCE_COLOR': ''}, fallback, fallback)
48+
check({'NO_COLOR': '1'}, fallback, False)
49+
check({'NO_COLOR': '0'}, fallback, False)
50+
check({'NO_COLOR': ''}, fallback, fallback)
51+
52+
check({'TERM': 'dumb', 'FORCE_COLOR': '1'}, False, True)
53+
check({'FORCE_COLOR': '1', 'NO_COLOR': '1'}, True, False)
54+
55+
for ignore_environment in False, True:
56+
# Simulate running with or without `-E`.
57+
flags = unittest.mock.MagicMock(ignore_environment=ignore_environment)
58+
with unittest.mock.patch("sys.flags", flags):
59+
check({'PYTHON_COLORS': '1'}, True, True)
60+
check({'PYTHON_COLORS': '1'}, False, not ignore_environment)
61+
check({'PYTHON_COLORS': '0'}, True, ignore_environment)
62+
check({'PYTHON_COLORS': '0'}, False, False)
63+
for fallback in False, True:
64+
check({'PYTHON_COLORS': 'x'}, fallback, fallback)
65+
check({'PYTHON_COLORS': ''}, fallback, fallback)
66+
67+
check({'TERM': 'dumb', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
68+
check({'NO_COLOR': '1', 'PYTHON_COLORS': '1'}, False, not ignore_environment)
69+
check({'FORCE_COLOR': '1', 'PYTHON_COLORS': '0'}, True, ignore_environment)
70+
71+
@unittest.skipUnless(sys.platform == "win32", "requires Windows")
72+
def test_colorized_detection_checks_on_windows(self):
73+
with (clear_env(),
74+
unittest.mock.patch("os.isatty") as isatty_mock,
75+
unittest.mock.patch("sys.stdout") as stdout_mock,
76+
supports_virtual_terminal() as vt_mock):
77+
stdout_mock.fileno.return_value = 1
78+
isatty_mock.return_value = True
79+
stdout_mock.isatty.return_value = True
80+
81+
vt_mock.return_value = True
82+
self.assertEqual(_colorize.can_colorize(), True)
83+
vt_mock.return_value = False
84+
self.assertEqual(_colorize.can_colorize(), False)
85+
import nt
86+
del nt._supports_virtual_terminal
87+
self.assertEqual(_colorize.can_colorize(), False)
88+
89+
def test_colorized_detection_checks_for_std_streams(self):
90+
with (clear_env(),
91+
unittest.mock.patch("os.isatty") as isatty_mock,
92+
unittest.mock.patch("sys.stdout") as stdout_mock,
93+
unittest.mock.patch("sys.stderr") as stderr_mock,
94+
supports_virtual_terminal()):
95+
stdout_mock.fileno.return_value = 1
96+
stderr_mock.fileno.side_effect = ZeroDivisionError
97+
stderr_mock.isatty.side_effect = ZeroDivisionError
98+
99+
isatty_mock.return_value = True
100+
stdout_mock.isatty.return_value = True
101+
self.assertEqual(_colorize.can_colorize(), True)
102+
103+
isatty_mock.return_value = False
104+
stdout_mock.isatty.return_value = False
105+
self.assertEqual(_colorize.can_colorize(), False)
106+
107+
def test_colorized_detection_checks_for_file(self):
108+
with clear_env(), supports_virtual_terminal():
109+
110+
with unittest.mock.patch("os.isatty") as isatty_mock:
111+
file = unittest.mock.MagicMock()
112+
file.fileno.return_value = 1
113+
isatty_mock.return_value = True
114+
self.assertEqual(_colorize.can_colorize(file=file), True)
115+
isatty_mock.return_value = False
116+
self.assertEqual(_colorize.can_colorize(file=file), False)
117+
118+
# No file.fileno.
119+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
120+
file = unittest.mock.MagicMock(spec=['isatty'])
121+
file.isatty.return_value = True
122+
self.assertEqual(_colorize.can_colorize(file=file), False)
123+
124+
# file.fileno() raises io.UnsupportedOperation.
125+
with unittest.mock.patch("os.isatty", side_effect=ZeroDivisionError):
126+
file = unittest.mock.MagicMock()
127+
file.fileno.side_effect = io.UnsupportedOperation
128+
file.isatty.return_value = True
129+
self.assertEqual(_colorize.can_colorize(file=file), True)
130+
file.isatty.return_value = False
131+
self.assertEqual(_colorize.can_colorize(file=file), False)
132+
133+
134+
if __name__ == "__main__":
135+
unittest.main()

‎vm/src/stdlib/nt.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,12 @@ pub(crate) mod module {
4343
|| attr & FileSystem::FILE_ATTRIBUTE_DIRECTORY != 0))
4444
}
4545

46+
#[pyfunction]
47+
pub(super) fn _supports_virtual_terminal() -> PyResult<bool> {
48+
// TODO: implement this
49+
Ok(true)
50+
}
51+
4652
#[derive(FromArgs)]
4753
pub(super) struct SymlinkArgs {
4854
src: OsPath,

0 commit comments

Comments
 (0)