Skip to content

Commit 19050af

Browse files
youknowoneyouknowone
authored
Merge pull request #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)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy