RustPython · youknowone · Feb 14, 2025 · Feb 14, 2025 · Feb 14, 2025
diff --git a/Lib/_colorize.py b/Lib/_colorize.py
@@ -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()
diff --git a/Lib/test/test__colorize.py b/Lib/test/test__colorize.py
@@ -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()
diff --git a/vm/src/stdlib/nt.rs b/vm/src/stdlib/nt.rs
@@ -43,6 +43,12 @@ pub(crate) mod module {
                || attr & FileSystem::FILE_ATTRIBUTE_DIRECTORY != 0))
    }

+    #[pyfunction]
+    pub(super) fn _supports_virtual_terminal() -> PyResult<bool> {
+        // TODO: implement this
+        Ok(true)
+    }
+
    #[derive(FromArgs)]
    pub(super) struct SymlinkArgs {
        src: OsPath,