fix: better way for already colorized

Signed-off-by: yihong0618 <zouzou0208@gmail.com>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

Author: yihong0618

Lib/_colorize.py

@@ -286,7 +286,10 @@ def decolor(text: str) -> str:
     return text
 
 
-def can_colorize(*, file: IO[str] | IO[bytes] | None = None) -> bool:
+def can_colorize(*, file: IO[str] | IO[bytes] | None = None, already_colorize=False) -> bool:
+    if already_colorize:
+        return True
+
     if file is None:
         file = sys.stdout
 
@@ -331,6 +334,7 @@ def get_theme(
     tty_file: IO[str] | IO[bytes] | None = None,
     force_color: bool = False,
     force_no_color: bool = False,
+    already_colorize: bool = False,
 ) -> Theme:
     """Returns the currently set theme, potentially in a zero-color variant.
 
@@ -344,7 +348,8 @@ def get_theme(
     environment (including environment variable state and console configuration
     on Windows) can also change in the course of the application life cycle.
     """
-    if force_color or (not force_no_color and can_colorize(file=tty_file)):
+    if force_color or (not force_no_color and
+                       can_colorize(file=tty_file, already_colorize=already_colorize)):
         return _theme
     return theme_no_color
 

Lib/_pyrepl/reader.py

@@ -343,7 +343,7 @@ def calc_screen(self) -> list[str]:
                 screeninfo.append((0, []))
             pos -= line_len + 1
             prompt, prompt_len = self.process_prompt(prompt)
-            chars, char_widths = disp_str(line, colors, offset)
+            chars, char_widths = disp_str(line, colors, offset, already_colorize=self.can_colorize)
             wrapcount = (sum(char_widths) + prompt_len) // self.console.width
             if wrapcount == 0 or not char_widths:
                 offset += line_len + 1  # Takes all of the line plus the newline
@@ -491,7 +491,7 @@ def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
             prompt = self.ps1
 
         if self.can_colorize:
-            t = THEME(force_color=self.can_colorize)
+            t = THEME(already_colorize=self.can_colorize)
             prompt = f"{t.prompt}{prompt}{t.reset}"
         return prompt
 

Lib/_pyrepl/unix_console.py

@@ -159,7 +159,10 @@ def __init__(
         self.pollob.register(self.input_fd, select.POLLIN)
         self.terminfo = terminfo.TermInfo(term or None)
         self.term = term
-        self.is_mac = platform.system() == "Darwin" and os.getenv("TERM_PROGRAM") == "Apple_Terminal"
+        self.is_mac = (
+            platform.system() == "Darwin"
+            and os.getenv("TERM_PROGRAM") == "Apple_Terminal"
+        )
 
         @overload
         def _my_getstr(cap: str, optional: Literal[False] = False) -> bytes: ...

Lib/_pyrepl/utils.py

@@ -266,6 +266,7 @@ def disp_str(
     colors: list[ColorSpan] | None = None,
     start_index: int = 0,
     force_color: bool = False,
+    already_colorize: bool = False,
 ) -> tuple[CharBuffer, CharWidths]:
     r"""Decompose the input buffer into a printable variant with applied colors.
 
@@ -306,7 +307,7 @@ def disp_str(
         # move past irrelevant spans
         colors.pop(0)
 
-    theme = THEME(force_color=force_color)
+    theme = THEME(force_color=force_color, already_colorize=already_colorize)
     pre_color = ""
     post_color = ""
     if colors and colors[0].span.start < start_index:

Lib/test/support/__init__.py

@@ -2898,7 +2898,7 @@ def force_color(color: bool):
     from .os_helper import EnvironmentVarGuard
 
     with (
-        swap_attr(_colorize, "can_colorize", lambda file=None: color),
+        swap_attr(_colorize, "can_colorize", lambda *, file=None, already_colorize=False: color),
         EnvironmentVarGuard() as env,
     ):
         env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS")

Lib/test/test_pyrepl/test_unix_console.py

@@ -305,17 +305,11 @@ def test_getheightwidth_with_invalid_environ(self, _os_write):
             self.assertIsInstance(console.getheightwidth(), tuple)
 
     @unittest.skipUnless(sys.platform == "darwin", "requires OS X")
-    def test_is_mac_with_invalid_environ(self, _os_write):
+    @unittest.skipUnless(sys.platform == "darwin", "requires OS X")
+    def test_restore_with_invalid_environ_on_macOS(self, _os_write):
         # gh-128636 for macOS
         console = UnixConsole(term="xterm")
         with os_helper.EnvironmentVarGuard() as env:
-            is_mac = os.getenv("TERM_PROGRAM") == "Apple_Terminal"
-            self.assertEqual(console.is_mac, is_mac)
             os.environ = []
-            console.prepare()  # Need to prepare before restore
-            try:
-                console.restore()  # This should not crash
-            except (KeyError, AttributeError, TypeError) as e:
-                self.fail(f"restore() failed with invalid environ: {e}")
-
-            self.assertEqual(console.is_mac, is_mac)
+            console.prepare()  # needed to call restore()
+            console.restore()  # this should succeed